@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8cdc371217
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 02:15 UTC

SM100 Cutlass MLA decode with unrestricted num_heads (< 128) for DeepSeek TP (#20769)

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

Alexander Matveev committed Jul 14, 2025 at 21:06 UTC 8cdc371217221bbddfe67a99a21a903609722a10
12 files changed +3283 -2
CMakeLists.txt
+2 -1
@@ -553,7 +553,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
553 cuda_archs_loose_intersection(MLA_ARCHS "10.0a" "${CUDA_ARCHS}")
554 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
555 set(SRCS
556 - "csrc/attention/mla/cutlass_mla_kernels.cu")
556 + "csrc/attention/mla/cutlass_mla_kernels.cu"
557 + "csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
558 set_gencode_flags_for_srcs(
559 SRCS "${SRCS}"
560 CUDA_ARCHS "${MLA_ARCHS}")
csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp new
+372
@@ -0,0 +1,372 @@
1 +/***************************************************************************************************
2 + * Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3 + * SPDX-License-Identifier: BSD-3-Clause
4 + *
5 + * Redistribution and use in source and binary forms, with or without
6 + * modification, are permitted provided that the following conditions are met:
7 + *
8 + * 1. Redistributions of source code must retain the above copyright notice,
9 + *this list of conditions and the following disclaimer.
10 + *
11 + * 2. Redistributions in binary form must reproduce the above copyright notice,
12 + * this list of conditions and the following disclaimer in the documentation
13 + * and/or other materials provided with the distribution.
14 + *
15 + * 3. Neither the name of the copyright holder nor the names of its
16 + * contributors may be used to endorse or promote products derived from
17 + * this software without specific prior written permission.
18 + *
19 + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20 + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21 + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22 + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
23 + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24 + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25 + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26 + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27 + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28 + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 + *POSSIBILITY OF SUCH DAMAGE.
30 + *
31 + **************************************************************************************************/
32 +/*
33 + * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
34 + * by Alcanderian JieXin Liang
35 + */
36 +
37 +/*!
38 + \file
39 + \brief An universal device layer for cutlass 3.x-style kernels.
40 +*/
41 +
42 +// clang-format off
43 +#pragma once
44 +
45 +// common
46 +#include "cutlass/cutlass.h"
47 +#include "cutlass/device_kernel.h"
48 +
49 +#if !defined(__CUDACC_RTC__)
50 +#include "cutlass/cluster_launch.hpp"
51 +#include "cutlass/trace.h"
52 +#endif // !defined(__CUDACC_RTC__)
53 +
54 +#include "../kernel/sm100_fmha_mla_tma_warpspecialized.hpp"
55 +#include "../kernel/sm100_fmha_mla_reduction.hpp"
56 +
57 +////////////////////////////////////////////////////////////////////////////////
58 +
59 +namespace cutlass::fmha::device {
60 +
61 +using namespace cute;
62 +using namespace cutlass::fmha::kernel;
63 +
64 +
65 +////////////////////////////////////////////////////////////////////////////////
66 +////////////////////////////// CUTLASS 3.x API /////////////////////////////////
67 +////////////////////////////////////////////////////////////////////////////////
68 +
69 +template<
70 + class Kernel_
71 +>
72 +class MLA {
73 +public:
74 +
75 + using Kernel = Kernel_;
76 +
77 + using ReductionKernel = cutlass::fmha::kernel::Sm100FmhaMlaReductionKernel<
78 + typename Kernel::ElementOut,
79 + typename Kernel::ElementAcc,
80 + typename Kernel::ElementAcc,
81 + Kernel::TileShapeH::value,
82 + Kernel::TileShapeL::value,
83 + 256 /*Max split*/
84 + >;
85 +
86 + /// Argument structure: User API
87 + using KernelArguments = typename Kernel::Arguments;
88 + using ReductionArguments = typename ReductionKernel::Arguments;
89 +
90 + using Arguments = KernelArguments;
91 +
92 + /// Argument structure: Kernel API
93 + using KernelParams = typename Kernel::Params;
94 + using ReductionParams = typename ReductionKernel::Params;
95 + struct Params {
96 + KernelParams fmha_params;
97 + ReductionParams reduction_params;
98 + };
99 +
100 +private:
101 +
102 + /// Kernel API parameters object
103 + Params params_;
104 +
105 + bool is_initialized(bool set = false) {
106 + static bool initialized = false;
107 + if (set) initialized = true;
108 + return initialized;
109 + }
110 +
111 + static ReductionArguments to_reduction_args(Arguments const& args) {
112 + auto [H, K, D, B] = args.problem_shape;
113 + return ReductionArguments{
114 + nullptr, args.epilogue.ptr_o, nullptr, args.epilogue.ptr_lse,
115 + args.mainloop.softmax_scale, B, args.split_kv, K, args.mainloop.ptr_seq,
116 + args.ptr_split_kv, Kernel::TileShapeS::value
117 + };
118 + }
119 +
120 +public:
121 +
122 + /// Access the Params structure
123 + Params const& params() const {
124 + return params_;
125 + }
126 +
127 + static void set_split_kv (KernelArguments& args) {
128 + // printf("set_split_kv start");
129 + if (args.split_kv >= 1) return;
130 + auto [H, K, D, B] = args.problem_shape;
131 + // std::cout << H << " " << K << " " << D << " " << B << "\n";
132 + int sm_count = args.hw_info.sm_count;
133 + // printf(" sm_count = %d\n", sm_count);
134 + int max_splits = ceil_div(K, 128);
135 + max_splits = min(16, max_splits);
136 + // printf(" max_splits = %d\n", max_splits);
137 + int sms_per_batch = max(1, sm_count / B);
138 + // printf(" sms_per_batch = %d\n", sms_per_batch);
139 + int split_heur = min(max_splits, sms_per_batch);
140 + int waves = ceil_div(B * split_heur, sm_count);
141 + int k_waves = ceil_div(max_splits, split_heur);
142 + int split_wave_aware = ceil_div(max_splits, k_waves);
143 + args.split_kv = split_wave_aware;
144 + // printf(" args.split_kv = %d\n", args.split_kv);
145 +
146 + }
147 +
148 + /// Determines whether the GEMM can execute the given problem.
149 + static Status
150 + can_implement(Arguments const& args) {
151 + if (! Kernel::can_implement(args)) {
152 + return Status::kInvalid;
153 + }
154 + if (! ReductionKernel::can_implement(to_reduction_args(args))) {
155 + return Status::kInvalid;
156 + }
157 + return Status::kSuccess;
158 + }
159 +
160 + /// Gets the workspace size
161 + static size_t
162 + get_workspace_size(Arguments const& args) {
163 + size_t workspace_bytes = 0;
164 + workspace_bytes += Kernel::get_workspace_size(args);
165 + workspace_bytes += ReductionKernel::get_workspace_size(to_reduction_args(args));
166 + return workspace_bytes;
167 + }
168 +
169 + /// Computes the maximum number of active blocks per multiprocessor
170 + static int maximum_active_blocks(int /* smem_capacity */ = -1) {
171 + CUTLASS_TRACE_HOST("MLA::maximum_active_blocks()");
172 + int max_active_blocks = -1;
173 + int smem_size = Kernel::SharedStorageSize;
174 +
175 + // first, account for dynamic smem capacity if needed
176 + cudaError_t result;
177 + if (smem_size >= (48 << 10)) {
178 + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
179 + result = cudaFuncSetAttribute(
180 + device_kernel<Kernel>,
181 + cudaFuncAttributeMaxDynamicSharedMemorySize,
182 + smem_size);
183 + if (cudaSuccess != result) {
184 + result = cudaGetLastError(); // to clear the error bit
185 + CUTLASS_TRACE_HOST(
186 + " cudaFuncSetAttribute() returned error: "
187 + << cudaGetErrorString(result));
188 + return -1;
189 + }
190 + }
191 +
192 + // query occupancy after setting smem size
193 + result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
194 + &max_active_blocks,
195 + device_kernel<Kernel>,
196 + Kernel::MaxThreadsPerBlock,
197 + smem_size);
198 +
199 + if (cudaSuccess != result) {
200 + result = cudaGetLastError(); // to clear the error bit
201 + CUTLASS_TRACE_HOST(
202 + " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: "
203 + << cudaGetErrorString(result));
204 + return -1;
205 + }
206 +
207 + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks);
208 + return max_active_blocks;
209 + }
210 +
211 + /// Initializes GEMM state from arguments.
212 + Status
213 + initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
214 + CUTLASS_TRACE_HOST("MLA::initialize() - workspace "
215 + << workspace << ", stream: " << (stream ? "non-null" : "null"));
216 +
217 + // Initialize the workspace
218 + Status status = Kernel::initialize_workspace(args, workspace, stream);
219 + if (status != Status::kSuccess) {
220 + return status;
221 + }
222 + status = ReductionKernel::initialize_workspace(to_reduction_args(args), workspace, stream);
223 + if (status != Status::kSuccess) {
224 + return status;
225 + }
226 + KernelParams kernel_params = Kernel::to_underlying_arguments(args, workspace);
227 +
228 + ReductionArguments reduction_args = to_reduction_args(args);
229 + if (reduction_args.split_kv > 1) {
230 + reduction_args.ptr_oaccum = kernel_params.epilogue.ptr_o_acc;
231 + reduction_args.ptr_lseaccum = kernel_params.epilogue.ptr_lse_acc;
232 + }
233 + ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace);
234 + // Initialize the Params structure
235 + params_ = Params {kernel_params, reduction_params};
236 +
237 + if (is_initialized()) return Status::kSuccess;
238 +
239 + // account for dynamic smem capacity if needed
240 + // no dynamic smem is needed for reduction kernel
241 + int smem_size = Kernel::SharedStorageSize;
242 + if (smem_size >= (48 << 10)) {
243 + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
244 + cudaError_t result = cudaFuncSetAttribute(
245 + device_kernel<Kernel>,
246 + cudaFuncAttributeMaxDynamicSharedMemorySize,
247 + smem_size);
248 + if (cudaSuccess != result) {
249 + result = cudaGetLastError(); // to clear the error bit
250 + CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
251 + return Status::kErrorInternal;
252 + }
253 + }
254 +
255 + is_initialized(true);
256 +
257 + return Status::kSuccess;
258 + }
259 +
260 + /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params.
261 + Status
262 + update(Arguments const& args, void* workspace = nullptr) {
263 + CUTLASS_TRACE_HOST("MLA()::update() - workspace: " << workspace);
264 +
265 + size_t workspace_bytes = get_workspace_size(args);
266 + if (workspace_bytes > 0 && nullptr == workspace) {
267 + return Status::kErrorWorkspaceNull;
268 + }
269 +
270 + auto fmha_params = Kernel::to_underlying_arguments(args, workspace);
271 +
272 + ReductionArguments reduction_args = to_reduction_args(args);
273 + if (reduction_args.split_kv > 1) {
274 + reduction_args.ptr_oaccum = fmha_params.epilogue.ptr_o_acc;
275 + reduction_args.ptr_lseaccum = fmha_params.epilogue.ptr_lse_acc;
276 + }
277 + ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace);
278 + // Initialize the Params structure
279 + params_ = Params {fmha_params, reduction_params};
280 +
281 + return Status::kSuccess;
282 + }
283 +
284 + /// Primary run() entry point API that is static allowing users to create and manage their own params.
285 + /// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
286 + static Status
287 + run(Params& params, cudaStream_t stream = nullptr) {
288 + CUTLASS_TRACE_HOST("MLA::run()");
289 + dim3 const block = Kernel::get_block_shape();
290 + dim3 const grid = Kernel::get_grid_shape(params.fmha_params);
291 +
292 + // configure smem size and carveout
293 + int smem_size = Kernel::SharedStorageSize;
294 +
295 + Status launch_result;
296 + // Use extended launch API only for mainloops that use it
297 + if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) {
298 + dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}),
299 + cute::size<1>(typename Kernel::ClusterShape{}),
300 + cute::size<2>(typename Kernel::ClusterShape{}));
301 + void const* kernel = (void const*) device_kernel<Kernel>;
302 + void* kernel_params[] = {&params.fmha_params};
303 + launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params);
304 + }
305 + else {
306 + launch_result = Status::kSuccess;
307 + device_kernel<Kernel><<<grid, block, smem_size, stream>>>(params.fmha_params);
308 + }
309 +
310 + cudaError_t result = cudaGetLastError();
311 + if (cudaSuccess != result or Status::kSuccess != launch_result) {
312 + //return Status::kSuccess;
313 + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
314 + return Status::kErrorInternal;
315 + }
316 + if (params.reduction_params.split_kv > 1) {
317 + // launch reduction kernel
318 + dim3 const block = ReductionKernel::get_block_shape();
319 + dim3 const grid = ReductionKernel::get_grid_shape(params.reduction_params);
320 + device_kernel<ReductionKernel><<<grid, block, 0, stream>>>(params.reduction_params);
321 + cudaError_t result = cudaGetLastError();
322 + if (cudaSuccess == result) {
323 + return Status::kSuccess;
324 + }
325 + else {
326 + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
327 + return Status::kErrorInternal;
328 + }
329 + }
330 + else {
331 + return Status::kSuccess;
332 + }
333 + }
334 +
335 + //
336 + // Non-static launch overloads that first create and set the internal params struct of this kernel handle.
337 + //
338 +
339 + /// Launches the kernel after first constructing Params internal state from supplied arguments.
340 + Status
341 + run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
342 + Status status = initialize(args, workspace, stream);
343 + if (Status::kSuccess == status) {
344 + status = run(params_, stream);
345 + }
346 + return status;
347 + }
348 +
349 + /// Launches the kernel after first constructing Params internal state from supplied arguments.
350 + Status
351 + operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
352 + return run(args, workspace, stream);
353 + }
354 +
355 + /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
356 + Status
357 + run(cudaStream_t stream = nullptr) {
358 + return run(params_, stream);
359 + }
360 +
361 + /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
362 + Status
363 + operator()(cudaStream_t stream = nullptr) {
364 + return run(params_, stream);
365 + }
366 +};
367 +
368 +////////////////////////////////////////////////////////////////////////////////
369 +
370 +} // namespace cutlass::fmha::device
371 +
372 +////////////////////////////////////////////////////////////////////////////////
csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp new
+203
@@ -0,0 +1,203 @@
1 +/***************************************************************************************************
2 + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
3 + *reserved. SPDX-License-Identifier: BSD-3-Clause
4 + *
5 + * Redistribution and use in source and binary forms, with or without
6 + * modification, are permitted provided that the following conditions are met:
7 + *
8 + * 1. Redistributions of source code must retain the above copyright notice,
9 + *this list of conditions and the following disclaimer.
10 + *
11 + * 2. Redistributions in binary form must reproduce the above copyright notice,
12 + * this list of conditions and the following disclaimer in the documentation
13 + * and/or other materials provided with the distribution.
14 + *
15 + * 3. Neither the name of the copyright holder nor the names of its
16 + * contributors may be used to endorse or promote products derived from
17 + * this software without specific prior written permission.
18 + *
19 + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20 + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21 + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22 + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
23 + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24 + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25 + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26 + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27 + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28 + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 + *POSSIBILITY OF SUCH DAMAGE.
30 + *
31 + **************************************************************************************************/
32 +/*
33 + * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
34 + * by Alcanderian JieXin Liang
35 + */
36 +
37 +// clang-format off
38 +#pragma once
39 +
40 +#include "cutlass/cutlass.h"
41 +#include "cutlass/arch/arch.h"
42 +#include "cute/tensor.hpp"
43 +
44 +namespace cutlass::fmha::kernel {
45 +
46 +using namespace cute;
47 +template<
48 + class ElementOut,
49 + class ElementAcc,
50 + class ElementScale,
51 + size_t kNumHeads,
52 + size_t kHeadDimLatent,
53 + int kMaxSplits
54 +>
55 +struct Sm100FmhaMlaReductionKernel {
56 +
57 + static const int SharedStorageSize = 0;
58 + static const int MaxThreadsPerBlock = 128;
59 + static const int MinBlocksPerMultiprocessor = 1;
60 +
61 + using ArchTag = cutlass::arch::Sm100;
62 +
63 + static_assert(kHeadDimLatent % MaxThreadsPerBlock == 0);
64 + struct Arguments {
65 + ElementAcc* ptr_oaccum = nullptr;
66 + ElementOut* ptr_o = nullptr;
67 + ElementAcc* ptr_lseaccum = nullptr;
68 + ElementAcc* ptr_lse = nullptr;
69 + ElementScale scale = 1.f;
70 + int num_batches = 0;
71 + int split_kv = -1;
72 + int dim_k = -1;
73 + int* ptr_seq = nullptr;
74 + int* ptr_split_kv = nullptr;
75 + int tile_shape_s = 128;
76 + };
77 + using Params = Arguments;
78 +
79 + static Params to_underlying_arguments(Arguments const& args, void* workspace) {
80 + return {args.ptr_oaccum, args.ptr_o, args.ptr_lseaccum, args.ptr_lse,
81 + args.scale, args.num_batches, args.split_kv, args.dim_k, args.ptr_seq,
82 + args.ptr_split_kv, args.tile_shape_s};
83 + }
84 +
85 + static size_t get_workspace_size(Arguments const& /*args*/) {
86 + return 0;
87 + }
88 +
89 + static Status initialize_workspace(
90 + Arguments const& /*args*/, void* /*ws*/, cudaStream_t /*stream*/) {
91 + return Status::kSuccess;
92 + }
93 +
94 + static dim3 get_grid_shape(Params const& params) {
95 + return dim3(kNumHeads, 1, params.num_batches);
96 + }
97 +
98 + static dim3 get_block_shape() {
99 + return dim3(MaxThreadsPerBlock, 1, 1);
100 + }
101 +
102 + static bool can_implement(Arguments const& args) {
103 + if (args.num_batches <= 0) return false;
104 + if (args.split_kv <= 0) return false;
105 + return true;
106 + }
107 +
108 + CUTLASS_DEVICE void operator() (Params const& params, char* smem_raw) {
109 + if (params.split_kv <= 1) return;
110 + auto blk_coord = make_coord(blockIdx.x, _0{}, blockIdx.z);
111 +
112 + __shared__ ElementAcc sLseScale[kMaxSplits];
113 + const size_t offset_lseaccum = get<0>(blk_coord) + kNumHeads * params.split_kv * get<2>(blk_coord);
114 + const size_t offset_lse = get<0>(blk_coord) + kNumHeads * get<2>(blk_coord);
115 +
116 + Tensor gLSEaccum = make_tensor(make_gmem_ptr(params.ptr_lseaccum + offset_lseaccum),
117 + make_shape(params.split_kv), Stride<Int<kNumHeads>>{});
118 +
119 + Tensor gLSE = make_tensor(make_gmem_ptr(params.ptr_lse + offset_lse),
120 + Shape<_1>{}, Stride<_1>{});
121 +
122 + auto dim_k = params.ptr_seq == nullptr ? params.dim_k : params.ptr_seq[get<2>(blk_coord)];
123 + auto local_split_kv = params.ptr_split_kv == nullptr ? params.split_kv : params.ptr_split_kv[get<2>(blk_coord)];
124 + auto k_tile_total = ceil_div(dim_k, params.tile_shape_s);
125 + auto k_tile_per_cta = ceil_div(k_tile_total, local_split_kv);
126 + local_split_kv = ceil_div(k_tile_total, k_tile_per_cta);
127 +
128 + int warp_idx = cutlass::canonical_warp_idx_sync();
129 + if (warp_idx == 0) {
130 + constexpr int kNLsePerThread = cute::ceil_div(kMaxSplits, 32);
131 +
132 + ElementAcc local_lse[kNLsePerThread];
133 +
134 + CUTLASS_PRAGMA_UNROLL
135 + for (int i = 0; i < kNLsePerThread; ++i) {
136 + const int split = i * 32 + threadIdx.x;
137 + local_lse[i] = split < local_split_kv ? gLSEaccum(split) : -std::numeric_limits<ElementAcc>::infinity();
138 + }
139 +
140 + ElementAcc lse_max = -std::numeric_limits<ElementAcc>::infinity();
141 + CUTLASS_PRAGMA_UNROLL
142 + for (int i = 0; i < kNLsePerThread; ++i) {
143 + lse_max = max(lse_max, local_lse[i]);
144 + }
145 + CUTLASS_PRAGMA_UNROLL
146 + for (int offset = 16; offset >= 1; offset /= 2) {
147 + lse_max = max(lse_max, __shfl_xor_sync(0xffffffff, lse_max, offset));
148 + }
149 + lse_max = lse_max == -std::numeric_limits<ElementAcc>::infinity() ? 0.0f : lse_max; // In case all local LSEs are -inf
150 + lse_max = __shfl_sync(0xffffffff, lse_max, 0);
151 +
152 + ElementAcc sum_lse = 0;
153 + CUTLASS_PRAGMA_UNROLL
154 + for (int i = 0; i < kNLsePerThread; ++i) {
155 + sum_lse = sum_lse + expf(local_lse[i] - lse_max);
156 + }
157 +
158 + CUTLASS_PRAGMA_UNROLL
159 + for (int offset = 16; offset >= 1; offset /= 2) {
160 + sum_lse = sum_lse + __shfl_xor_sync(0xffffffff, sum_lse, offset);
161 + }
162 +
163 + sum_lse = __shfl_sync(0xffffffff, sum_lse, 0);
164 +
165 + ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? std::numeric_limits<ElementAcc>::infinity() : logf(sum_lse) + lse_max;
166 + if (threadIdx.x == 0 and params.ptr_lse != nullptr) {
167 + gLSE(0) = global_lse;
168 + }
169 +
170 + CUTLASS_PRAGMA_UNROLL
171 + for (int i = 0; i < kNLsePerThread; ++i) {
172 + const int split = i * 32 + threadIdx.x;
173 + if (split < local_split_kv) {
174 + sLseScale[split] = expf(local_lse[i] - global_lse);
175 + }
176 + }
177 + }
178 + __syncthreads();
179 +
180 + constexpr int Elements = kHeadDimLatent / MaxThreadsPerBlock;
181 + const size_t offset_oaccum = kHeadDimLatent * params.split_kv * (get<0>(blk_coord) + kNumHeads * get<2>(blk_coord));
182 + Tensor gOaccum = make_tensor(make_gmem_ptr(params.ptr_oaccum + offset_oaccum),
183 + Shape<Int<kHeadDimLatent>>{}, Stride<_1>{});
184 + ElementAcc local_val[Elements] = {0};
185 + for (int split = 0; split < local_split_kv; ++split) {
186 + ElementAcc lse_scale = sLseScale[split];
187 + CUTLASS_PRAGMA_UNROLL
188 + for(int i = 0; i < Elements; ++i) {
189 + local_val[i] += lse_scale * gOaccum(threadIdx.x + MaxThreadsPerBlock * i);
190 + }
191 + gOaccum.data() = gOaccum.data() + kHeadDimLatent;
192 + }
193 + auto ptr_o_local = params.ptr_o + (get<0>(blk_coord) + get<2>(blk_coord) * kNumHeads) * kHeadDimLatent;
194 + Tensor gO = make_tensor(make_gmem_ptr(ptr_o_local), Shape<Int<kHeadDimLatent>>{}, Stride<_1>{});
195 +
196 + CUTLASS_PRAGMA_UNROLL
197 + for(int i = 0; i < Elements; ++i) {
198 + gO(threadIdx.x + MaxThreadsPerBlock * i) = static_cast<ElementOut>(local_val[i]);
199 + }
200 + }
201 +};
202 +
203 +} // namespace cutlass::fmha::kernel
csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp new
+2023
@@ -0,0 +1,2023 @@
1 +/***************************************************************************************************
2 + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
3 + *reserved. SPDX-License-Identifier: BSD-3-Clause
4 + *
5 + * Redistribution and use in source and binary forms, with or without
6 + * modification, are permitted provided that the following conditions are met:
7 + *
8 + * 1. Redistributions of source code must retain the above copyright notice,
9 + *this list of conditions and the following disclaimer.
10 + *
11 + * 2. Redistributions in binary form must reproduce the above copyright notice,
12 + * this list of conditions and the following disclaimer in the documentation
13 + * and/or other materials provided with the distribution.
14 + *
15 + * 3. Neither the name of the copyright holder nor the names of its
16 + * contributors may be used to endorse or promote products derived from
17 + * this software without specific prior written permission.
18 + *
19 + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20 + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21 + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22 + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
23 + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24 + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25 + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26 + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27 + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28 + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 + *POSSIBILITY OF SUCH DAMAGE.
30 + *
31 + **************************************************************************************************/
32 +/*
33 + * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
34 + * by Alcanderian JieXin Liang
35 + */
36 +
37 +// clang-format off
38 +#pragma once
39 +
40 +#include "cutlass/cutlass.h"
41 +
42 +#include "cute/tensor.hpp"
43 +#include "cute/arch/simd_sm100.hpp"
44 +
45 +#include "cutlass/arch/arch.h"
46 +#include "cutlass/arch/memory_sm80.h"
47 +#include "cutlass/epilogue/thread/linear_combination.h"
48 +#include "cutlass/gemm/collective/collective_builder.hpp"
49 +
50 +#include "gather_tensor.hpp" // from examples/common
51 +#include "common/pow_2.hpp"
52 +
53 +namespace cutlass::fmha::kernel {
54 +
55 +using namespace cute;
56 +
57 +template<
58 + class TileShape,
59 + class Element_,
60 + class ElementAcc_,
61 + class ElementOut_,
62 + class ElementLSE_,
63 + class TileScheduler,
64 +#ifdef CPASYNC
65 + bool kIsCpAsync = true
66 +#else
67 + bool kIsCpAsync = false
68 +#endif
69 +>
70 +struct Sm100FmhaMlaKernelTmaWarpspecialized {
71 +
72 + using Element = Element_;
73 + using ElementAcc = ElementAcc_;
74 + using ElementOut = ElementOut_;
75 + using ElementLSE = ElementLSE_;
76 +
77 + // only 2Sm mode is supported
78 + static const bool kIs2Sm = true;
79 + static const int MaxThreadsPerBlock = 256;
80 + static const int MinBlocksPerMultiprocessor = 1;
81 + static const int TotalSNum = 2;
82 + static const int TotalPNum = 2;
83 + using ArchTag = cutlass::arch::Sm100;
84 +
85 + using ClusterShape = cute::conditional_t<kIs2Sm, Shape<_2, _1, _1>, Shape<_1, _1, _1>>;
86 +
87 + using TileShapeH = tuple_element_t<0, TileShape>;
88 + using TileShapeS = tuple_element_t<1, TileShape>;
89 + using TileShapeD = tuple_element_t<2, TileShape>;
90 +
91 + using TileShapeL = tuple_element_t<0, TileShapeD>;
92 + using TileShapeR = tuple_element_t<1, TileShapeD>;
93 + static_assert(TileShapeL{} % TileShapeR{} == 0, "Rope head dim must divide latent head dim");
94 +
95 + using ProblemShape = Shape<TileShapeH, int, TileShapeD, int>;
96 + using TensorStride = Stride<int64_t, _1, int64_t>;
97 + using TmemAllocator = cute::conditional_t<kIs2Sm, cute::TMEM::Allocator2Sm, cute::TMEM::Allocator1Sm>;
98 +
99 + static_assert(TileShapeH{} == 128);
100 + static const int kWarpsInN = kIs2Sm ? 2 : 1;
101 +
102 + static const int kNumComputeWarps = 4;
103 + static const int kNumLoadWarps = kIsCpAsync ? 2 : 1;
104 +
105 + enum class WarpRole {
106 + kMma = 0x1, kLoad = 0x2, kCompute = 0x3, kLoadPageTable = 0x4, kEmpty=0x0
107 + };
108 +
109 + static const long long unsigned int kWarpAssignment = kIsCpAsync ? 0x4221'3333ull : 0x0021'3333ull;
110 +
111 + static CUTLASS_DEVICE WarpRole warp_idx_to_role(int warp_idx) {
112 + return static_cast<WarpRole>((kWarpAssignment >> (4 * warp_idx)) & 0xF);
113 + }
114 +
115 + static const int Alignment = 128 / sizeof_bits_v<Element>;
116 + static const int AlignmentOut = 128 / sizeof_bits_v<ElementOut>;
117 +
118 + using TileShapeQK = Shape<TileShapeH, TileShapeS, decltype(TileShapeR{} / _1{})>;
119 + static const int StagesQK = 24 / sizeof(Element); // free parameter
120 + static const int IterationsQKLatent = decltype(TileShapeL{} / get<2>(TileShapeQK{}))::value;
121 + static const int IterationsQKRope = decltype(TileShapeR{} / get<2>(TileShapeQK{}))::value;
122 + static const int IterationsQK = IterationsQKLatent + IterationsQKRope;
123 +
124 + using Schedule = cute::conditional_t<kIs2Sm, cutlass::gemm::KernelTmaWarpSpecialized2SmSm100, cutlass::gemm::KernelTmaWarpSpecialized1SmSm100>;
125 + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder<
126 + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
127 + Element, TensorStride, Alignment,
128 + Element, TensorStride, Alignment,
129 + ElementAcc,
130 + TileShapeQK, ClusterShape, cutlass::gemm::collective::StageCount<StagesQK>,
131 + Schedule>::CollectiveOp;
132 + using TiledMmaQK = typename CollectiveMmaQK::TiledMma;
133 + using CtaShapeQK = typename CollectiveMmaQK::CtaShape_MNK;
134 +
135 + // chosen for unified smem staging between K and V
136 + using TileShapePV = Shape<TileShapeH, _256, _32>;
137 + using TransposeTensorStride = decltype(select<1,0,2>(TensorStride{}));
138 + static const int StagesPV = StagesQK; // not sure why, but must be at least two. check pipes
139 + static const int IterationsPV_K = decltype(TileShapeS{} / get<2>(TileShapePV{}))::value;
140 + static const int IterationsPV_N = decltype(TileShapeL{} / get<1>(TileShapePV{}))::value;
141 +
142 + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder<
143 + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
144 + Element, TensorStride, Alignment,
145 + Element, TransposeTensorStride, Alignment,
146 + ElementAcc,
147 + TileShapePV, ClusterShape, cutlass::gemm::collective::StageCount<StagesPV>,
148 + Schedule>::CollectiveOp;
149 + using CtaShapePV = typename CollectiveMmaPV::CtaShape_MNK;
150 + static_assert(std::is_same_v<TransposeTensorStride, typename CollectiveMmaPV::StrideB>);
151 +
152 + using TiledMmaPV = typename CollectiveMmaPV::TiledMma;
153 +
154 + using AtomThrShapeMNK = typename CollectiveMmaQK::AtomThrShapeMNK;
155 + static_assert(typename CollectiveMmaQK::AtomThrShapeMNK{} == typename CollectiveMmaPV::AtomThrShapeMNK{}, "schedule must match");
156 +
157 + static const int StagesPageTable = kIsCpAsync ? StagesPV : 1;
158 +
159 + // pipelines from load to mma, PipelineTmaUmmaAsync, stages tbd
160 + // use expect_tx for Q load
161 + using PipelineLoadQK = cute::conditional_t<kIsCpAsync, PipelineUmmaConsumerAsync<StagesQK, AtomThrShapeMNK>, PipelineTmaUmmaAsync<StagesQK, ClusterShape, AtomThrShapeMNK>>;
162 + using PipelineLoadPV = PipelineLoadQK;
163 + // pipeline from mma (Q@K) to softmax, PipelineUmmaAsync, 2 stages
164 + using PipelineS = PipelineUmmaAsync<TotalSNum, AtomThrShapeMNK>;
165 + // pipeline from softmax (P) to mma (bmm2), PipelineUmmaAsync, 2 stages
166 + using PipelineP = PipelineUmmaConsumerAsync<TotalPNum, AtomThrShapeMNK>;
167 + // pipeline from mma to softmax (for rescale), PipelineUmmaAsync, 1 stage
168 + using PipelineO = PipelineUmmaAsync<1, AtomThrShapeMNK>;
169 +
170 + using PipelinePT = PipelineAsync<StagesPageTable>;
171 +
172 + struct PipelineStorage {
173 + alignas(16) typename PipelineLoadQK::SharedStorage load_qk;
174 + alignas(16) typename PipelineS::SharedStorage mma_s;
175 + alignas(16) typename PipelineP::SharedStorage p_mma;
176 + alignas(16) typename PipelineO::SharedStorage mma_o;
177 + alignas(16) typename PipelinePT::SharedStorage load_page_table;
178 + };
179 +
180 + template<class Layout, class Stages = _1>
181 + static CUTE_DEVICE constexpr auto unstageSmemLayout(Layout const& layout, Stages stages = {}) {
182 + return composition(layout, make_tuple(_, _, _, make_layout(stages)));
183 + }
184 +
185 + using SmemLayoutQ = decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutA{}, Int<IterationsQK>{}));
186 + using SmemLayoutKC = typename CollectiveMmaQK::SmemLayoutB;
187 + using SmemLayoutVC = typename CollectiveMmaPV::SmemLayoutB;
188 + using SmemLayoutP = decltype(unstageSmemLayout(typename CollectiveMmaPV::SmemLayoutA{}, make_shape(Int<IterationsPV_K>{}, _2{})));
189 +
190 + static const int kBytesLoadQ = size(AtomThrShapeMNK{}) * cutlass::bits_to_bytes(cosize(take<0,3>(SmemLayoutQ{})) * cute::sizeof_bits_v<Element>);
191 + static const int kBytesLoadKC = size(AtomThrShapeMNK{}) * cutlass::bits_to_bytes(cosize(take<0,3>(SmemLayoutKC{})) * cute::sizeof_bits_v<Element>);
192 + static const int kBytesLoadVC = size(AtomThrShapeMNK{}) * cutlass::bits_to_bytes(cosize(take<0,3>(SmemLayoutVC{})) * cute::sizeof_bits_v<Element>);
193 + // pre-condition for overlapped smem staging
194 + static_assert(kBytesLoadKC == kBytesLoadVC);
195 + static_assert(StagesQK == StagesPV);
196 +
197 + static const int kTransactionsBytesLoadQK = kBytesLoadKC;
198 + static const int kTransactionsBytesLoadExtraQ = kBytesLoadQ;
199 + static const int kTransactionsBytesLoadPV = kBytesLoadVC;
200 +
201 + static const int kNamedBarrierExchange = (int) cutlass::arch::ReservedNamedBarriers::TransformBarrier;
202 + // This Named Barrier is introduced to solve Q tile loading overwritten issue when enable persistent
203 + // tile scheduler for FP8 MLA.
204 + static const int kNamedBarrierEpilogue = (int) cutlass::arch::ReservedNamedBarriers::EpilogueBarrier;
205 + //
206 + static const int kNamedBarrierTmemDealloc = (int) cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier;
207 +
208 + enum class TmemAllocation : uint32_t {
209 + kSizeS = TileShapeS::value / kWarpsInN,
210 + // Overall
211 + kSizeO = TileShapeL::value / kWarpsInN,
212 + // Between accumulators we loop over
213 + kSizeAccO = decltype(get<1>(TileShapePV{}))::value / kWarpsInN,
214 + kNumS = TotalSNum,
215 + kNumP = TotalPNum,
216 + kNumO = 1,
217 + kS0 = 0,
218 + kS1 = kS0 + kSizeS,
219 + kO0 = kS1 + kSizeS,
220 + kTotal = kO0 + kSizeO
221 + };
222 +
223 + static_assert(static_cast<int>(TmemAllocation::kTotal) <= TmemAllocator::Sm100TmemCapacityColumns, "using too much tmem");
224 +
225 + struct TensorStorage {
226 + // to communicate max and row_sum
227 + cute::array<ElementAcc, kNumComputeWarps * cutlass::NumThreadsPerWarp> smem_exchange;
228 + cute::array<int, StagesPageTable * TileShapeS::value> smem_page_table;
229 + alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutQ>> smem_q;
230 + union {
231 + alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutKC>> smem_kc;
232 + alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutVC>> smem_vc;
233 + };
234 + alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutP>> smem_p;
235 + };
236 +
237 + struct SharedStorage {
238 + PipelineStorage pipelines;
239 + TensorStorage tensors;
240 + uint32_t tmem_base_ptr;
241 + };
242 +
243 + static const int SharedStorageSize = sizeof(SharedStorage);
244 + static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "using too much smem");
245 +
246 + struct MainloopArguments {
247 + ElementAcc softmax_scale;
248 +
249 + // all tensors strides are (num_heads or seqlen, head_dim, batch)
250 + // head_dim stride is always 1
251 + Element* ptr_q_latent;
252 + TensorStride stride_q_latent;
253 + Element* ptr_q_rope;
254 + TensorStride stride_q_rope;
255 +
256 + Element* ptr_c_latent;
257 + TensorStride stride_c_latent;
258 + Element* ptr_k_rope;
259 + TensorStride stride_k_rope;
260 +
261 + // for paged attention, we interpret what was previously [batch, seqlen]
262 + // as [page_count, page_size], and index according to page_table
263 + int* ptr_seq = nullptr;
264 + int* ptr_page_table = nullptr;
265 + // page table is [batch, seqlen or similar]
266 + Stride<_1, int> stride_page_table = {};
267 + int page_count = 0;
268 + int page_size = TileShapeS{}; // powers of two if kIsCpAsync, otherwise TileShapeS
269 + };
270 +
271 + struct EpilogueArguments {
272 + ElementOut* ptr_o = nullptr;
273 + TensorStride stride_o;
274 + ElementLSE* ptr_lse = nullptr;
275 + Stride<_1, int> stride_lse;
276 + ElementAcc output_scale = 1.0f;
277 + };
278 +
279 + struct Arguments {
280 + // (num_heads=128, seqlen, (d_latent=512, d_rope=64), batch_count)
281 + // for paged attention, seqlen is max seqlen
282 + ProblemShape problem_shape;
283 + MainloopArguments mainloop;
284 + EpilogueArguments epilogue;
285 + KernelHardwareInfo hw_info;
286 + int split_kv = -1;
287 + int* ptr_split_kv = nullptr;
288 + };
289 +
290 + using TmaLoadQLatent = typename CollectiveMmaQK::Params::TMA_A;
291 + using TmaLoadQRope = typename CollectiveMmaQK::Params::TMA_A;
292 + using TmaLoadCLatent = typename CollectiveMmaQK::Params::TMA_B;
293 + using TmaLoadKRope = typename CollectiveMmaQK::Params::TMA_B;
294 + using TmaLoadCLatentTranspose = typename CollectiveMmaPV::Params::TMA_B;
295 +
296 + struct MainloopParams {
297 + TmaLoadQLatent tma_load_q_latent;
298 + TmaLoadQRope tma_load_q_rope;
299 + TmaLoadCLatent tma_load_c_latent;
300 + TmaLoadKRope tma_load_k_rope;
301 + TmaLoadCLatentTranspose tma_load_c_latent_transpose;
302 + };
303 +
304 + struct EpilogueParams {
305 + ElementOut* ptr_o = nullptr;
306 + ElementAcc* ptr_o_acc = nullptr;
307 + TensorStride stride_o;
308 + TensorStride stride_o_acc;
309 + ElementLSE* ptr_lse = nullptr;
310 + ElementLSE* ptr_lse_acc = nullptr;
311 + Stride<_1, int> stride_lse;
312 + Stride<_1, int> stride_lse_acc;
313 + ElementAcc output_scale = 1.0f;
314 + };
315 +
316 + struct Params {
317 + ProblemShape problem_shape;
318 + MainloopArguments mainloop;
319 + EpilogueParams epilogue;
320 + MainloopParams mainloop_params;
321 + typename TileScheduler::Params tile_scheduler;
322 + int split_kv = -1;
323 + int* ptr_split_kv = nullptr;
324 + };
325 +
326 + static Params to_underlying_arguments(Arguments const& args, void* workspace) {
327 + //workspace = nullptr; // let's get an error if one of these needs workspace
328 +
329 + auto [H, K, D, B] = args.problem_shape;
330 + auto [L, R] = D;
331 +
332 + int paged_B = B;
333 + int paged_K = K;
334 + if (args.mainloop.ptr_page_table != nullptr) {
335 + paged_B = args.mainloop.page_count;
336 + paged_K = args.mainloop.page_size;
337 + }
338 +
339 + auto params_qk_latent = CollectiveMmaQK::to_underlying_arguments(
340 + make_shape(H, K, L, B),
341 + typename CollectiveMmaQK::Arguments {
342 + args.mainloop.ptr_q_latent, args.mainloop.stride_q_latent,
343 + args.mainloop.ptr_c_latent, args.mainloop.stride_c_latent,
344 + }, nullptr);
345 +
346 + auto params_qk_latent_paged = CollectiveMmaQK::to_underlying_arguments(
347 + make_shape(H, paged_K, L, paged_B),
348 + typename CollectiveMmaQK::Arguments {
349 + args.mainloop.ptr_q_latent, args.mainloop.stride_q_latent,
350 + args.mainloop.ptr_c_latent, args.mainloop.stride_c_latent,
351 + }, nullptr);
352 +
353 + auto params_qk_rope = CollectiveMmaQK::to_underlying_arguments(
354 + make_shape(H, K, R, B),
355 + typename CollectiveMmaQK::Arguments {
356 + args.mainloop.ptr_q_rope, args.mainloop.stride_q_rope,
357 + args.mainloop.ptr_k_rope, args.mainloop.stride_k_rope,
358 + }, nullptr);
359 +
360 + auto params_qk_rope_paged = CollectiveMmaQK::to_underlying_arguments(
361 + make_shape(H, paged_K, R, paged_B),
362 + typename CollectiveMmaQK::Arguments {
363 + args.mainloop.ptr_q_rope, args.mainloop.stride_q_rope,
364 + args.mainloop.ptr_k_rope, args.mainloop.stride_k_rope,
365 + }, nullptr);
366 +
367 +
368 + auto stride_c_latent_transpose = select<1,0,2>(args.mainloop.stride_c_latent);
369 + auto params_pv_latent = CollectiveMmaPV::to_underlying_arguments(
370 + make_shape(H, L, paged_K, paged_B),
371 + typename CollectiveMmaPV::Arguments {
372 + args.mainloop.ptr_q_latent, args.mainloop.stride_q_latent, // dummy, never used
373 + args.mainloop.ptr_c_latent, stride_c_latent_transpose,
374 + }, nullptr);
375 +
376 + MainloopParams mainloop_params {
377 + params_qk_latent.tma_load_a,
378 + params_qk_rope.tma_load_a,
379 + params_qk_latent_paged.tma_load_b,
380 + params_qk_rope_paged.tma_load_b,
381 + params_pv_latent.tma_load_b
382 + };
383 +
384 + EpilogueParams epilogue_params;
385 +
386 + epilogue_params.ptr_o = args.epilogue.ptr_o;
387 + epilogue_params.stride_o = args.epilogue.stride_o;
388 + epilogue_params.ptr_lse = args.epilogue.ptr_lse;
389 + epilogue_params.stride_lse = args.epilogue.stride_lse;
390 + epilogue_params.output_scale = args.epilogue.output_scale;
391 +
392 + if (args.split_kv > 1) {
393 + ElementAcc* ptr_o_acc = reinterpret_cast<ElementAcc*>(workspace);
394 + ElementLSE* ptr_lse_acc = reinterpret_cast<ElementLSE*>(ptr_o_acc + H * L * args.split_kv * B);
395 + epilogue_params.ptr_o_acc = ptr_o_acc;
396 + epilogue_params.ptr_lse_acc = ptr_lse_acc;
397 +
398 + epilogue_params.stride_o_acc = make_tuple(static_cast<int64_t>(0 + L) * args.split_kv, _1{}, static_cast<int64_t>(0 + H * L) * args.split_kv);
399 + epilogue_params.stride_lse_acc = make_tuple(_1{}, (0 + H) * args.split_kv);
400 + }
401 +
402 + return {args.problem_shape, args.mainloop, epilogue_params, mainloop_params,
403 + TileScheduler::to_underlying_arguments(args.problem_shape, args.hw_info, ClusterShape{}, args.split_kv), args.split_kv, args.ptr_split_kv};
404 + }
405 +
406 + static size_t get_workspace_size(Arguments const& args) {
407 + ProblemShape problem_shape = args.problem_shape;
408 + auto [H, K, D, B] = problem_shape;
409 + auto [D_latent, D_rope] = D;
410 + auto split_kv = args.split_kv;
411 + return (sizeof(ElementAcc) * D_latent + sizeof(ElementLSE)) * H * split_kv * B;
412 + }
413 + static Status initialize_workspace(
414 + Arguments const& /*args*/, void* /*ws*/, cudaStream_t /*stream*/) {
415 + return Status::kSuccess;
416 + }
417 +
418 + static dim3 get_grid_shape(Params const& params) {
419 + return TileScheduler::get_grid_shape(params.tile_scheduler);
420 + }
421 +
422 + static dim3 get_block_shape() {
423 + dim3 block(MaxThreadsPerBlock, 1, 1);
424 + return block;
425 + }
426 +
427 + static bool can_implement(Arguments const& args) {
428 + if (kIsCpAsync) {
429 + if ((args.mainloop.page_size & (args.mainloop.page_size - 1)) != 0) {
430 + return false;
431 + }
432 + if (args.mainloop.page_size > TileShapeS{}) {
433 + return false;
434 + }
435 + }
436 + else {
437 + if (args.mainloop.ptr_page_table != nullptr && args.mainloop.page_size != TileShapeS{}) {
438 + return false;
439 + }
440 + }
441 + if (get<0>(args.problem_shape) != 128) {
442 + return false;
443 + }
444 + if (get<1>(args.problem_shape) <= 0) {
445 + return false;
446 + }
447 + if (args.split_kv <= 0) {
448 + return false;
449 + }
450 + return true;
451 + }
452 +
453 +
454 + CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) {
455 +
456 + TileScheduler tile_scheduler(params.tile_scheduler);
457 +
458 + int warp_idx = cutlass::canonical_warp_idx_sync();
459 + auto role = warp_idx_to_role(warp_idx);
460 + uint32_t lane_predicate = cute::elect_one_sync();
461 +
462 + uint32_t cta_rank_in_cluster = cute::block_rank_in_cluster();
463 + int cta_coord_v = cta_rank_in_cluster % size<0>(AtomThrShapeMNK{});
464 + bool is_mma_leader_cta = cta_coord_v == 0;
465 +
466 + if (role == WarpRole::kLoad && lane_predicate && ! kIsCpAsync) {
467 + prefetch_tma_descriptor(params.mainloop_params.tma_load_q_latent.get_tma_descriptor());
468 + prefetch_tma_descriptor(params.mainloop_params.tma_load_c_latent.get_tma_descriptor());
469 + prefetch_tma_descriptor(params.mainloop_params.tma_load_q_rope.get_tma_descriptor());
470 + prefetch_tma_descriptor(params.mainloop_params.tma_load_k_rope.get_tma_descriptor());
471 + prefetch_tma_descriptor(params.mainloop_params.tma_load_c_latent_transpose.get_tma_descriptor());
472 + }
473 + SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_raw);
474 +
475 + typename PipelineLoadQK::Params pipeline_load_qk_params;
476 + if (role == WarpRole::kLoad) {
477 + pipeline_load_qk_params.role = PipelineLoadQK::ThreadCategory::Producer;
478 + }
479 + if (role == WarpRole::kMma) {
480 + pipeline_load_qk_params.role = PipelineLoadQK::ThreadCategory::Consumer;
481 + }
482 + if constexpr (kIsCpAsync) {
483 + // we can make our life easier by unconditionally loading blocks
484 + // since we know it'll always be legal
485 + pipeline_load_qk_params.producer_arv_count = kNumLoadWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{});
486 + }
487 + else {
488 + pipeline_load_qk_params.is_leader = lane_predicate && (role == WarpRole::kLoad) && is_mma_leader_cta;
489 + pipeline_load_qk_params.transaction_bytes = kTransactionsBytesLoadQK;
490 + }
491 + pipeline_load_qk_params.initializing_warp = 0;
492 + PipelineLoadQK pipeline_load_qk(shared_storage.pipelines.load_qk, pipeline_load_qk_params,
493 + ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{});
494 +
495 + typename PipelineS::Params pipeline_mma_s_params;
496 + if (role == WarpRole::kMma) {
497 + pipeline_mma_s_params.role = PipelineS::ThreadCategory::Producer;
498 + }
499 + if (role == WarpRole::kCompute) {
500 + pipeline_mma_s_params.role = PipelineS::ThreadCategory::Consumer;
501 + }
502 + pipeline_mma_s_params.consumer_arv_count = kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{});
503 + pipeline_mma_s_params.initializing_warp = 1;
504 + PipelineS pipeline_mma_s(
505 + shared_storage.pipelines.mma_s,
506 + pipeline_mma_s_params,
507 + ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{});
508 +
509 + typename PipelineP::Params pipeline_p_mma_params;
510 + if (role == WarpRole::kMma) {
511 + pipeline_p_mma_params.role = PipelineP::ThreadCategory::Consumer;
512 + }
513 + if (role == WarpRole::kCompute) {
514 + pipeline_p_mma_params.role = PipelineP::ThreadCategory::Producer;
515 + }
516 + pipeline_p_mma_params.producer_arv_count = kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{});
517 + pipeline_p_mma_params.consumer_arv_count = 1;
518 + pipeline_p_mma_params.initializing_warp = 2;
519 + PipelineP pipeline_p_mma(
520 + shared_storage.pipelines.p_mma,
521 + pipeline_p_mma_params,
522 + ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{});
523 +
524 + typename PipelineO::Params pipeline_mma_o_params;
525 + if (role == WarpRole::kMma) {
526 + pipeline_mma_o_params.role = PipelineO::ThreadCategory::Producer;
527 + }
528 + if (role == WarpRole::kCompute) {
529 + pipeline_mma_o_params.role = PipelineO::ThreadCategory::Consumer;
530 + }
531 + pipeline_mma_o_params.consumer_arv_count = kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{});
532 + pipeline_mma_o_params.initializing_warp = 3;
533 + PipelineO pipeline_mma_o(
534 + shared_storage.pipelines.mma_o,
535 + pipeline_mma_o_params,
536 + ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{});
537 +
538 + typename PipelinePT::Params pipeline_pt_params;
539 + if (role == WarpRole::kLoad) {
540 + pipeline_pt_params.role = PipelinePT::ThreadCategory::Consumer;
541 + }
542 + if (role == WarpRole::kLoadPageTable) {
543 + pipeline_pt_params.role = PipelinePT::ThreadCategory::Producer;
544 + }
545 + pipeline_pt_params.consumer_arv_count = kNumLoadWarps * cutlass::NumThreadsPerWarp;
546 + pipeline_pt_params.producer_arv_count = cutlass::NumThreadsPerWarp;
547 + pipeline_pt_params.initializing_warp = 4;
548 + PipelinePT pipeline_page_table(
549 + shared_storage.pipelines.load_page_table,
550 + pipeline_pt_params);
551 +
552 + TmemAllocator tmem_allocator;
553 +
554 + pipeline_init_arrive_relaxed(size(ClusterShape{}));
555 +
556 + pipeline_load_qk.init_masks(ClusterShape{}); // do we need an update here for 2Sm?
557 + pipeline_mma_s.init_masks(ClusterShape{});
558 + pipeline_p_mma.init_masks(ClusterShape{});
559 + pipeline_mma_o.init_masks(ClusterShape{});
560 +
561 + typename PipelineLoadQK::PipelineState pipeline_load_qk_consumer_state;
562 + typename PipelineLoadQK::PipelineState pipeline_load_qk_producer_state = cutlass::make_producer_start_state<PipelineLoadQK>();
563 +
564 + typename PipelineS::PipelineState pipeline_mma_s_consumer_state;
565 + typename PipelineS::PipelineState pipeline_mma_s_producer_state = cutlass::make_producer_start_state<PipelineS>();
566 +
567 + typename PipelineP::PipelineState pipeline_p_mma_consumer_state;
568 + typename PipelineP::PipelineState pipeline_p_mma_producer_state = cutlass::make_producer_start_state<PipelineP>();
569 +
570 + typename PipelineO::PipelineState pipeline_mma_o_consumer_state;
571 + typename PipelineO::PipelineState pipeline_mma_o_producer_state = cutlass::make_producer_start_state<PipelineO>();
572 +
573 + typename PipelinePT::PipelineState pipeline_pt_consumer_state;
574 + typename PipelinePT::PipelineState pipeline_pt_producer_state = cutlass::make_producer_start_state<PipelinePT>();
575 +
576 + pipeline_init_wait(size(ClusterShape{}));
577 +
578 + if (role == WarpRole::kLoadPageTable) {
579 + CUTLASS_PRAGMA_NO_UNROLL
580 + for (; tile_scheduler.is_valid(); ++tile_scheduler) {
581 + auto blk_coord = tile_scheduler.get_block_coord();
582 + auto problem_shape = params.problem_shape;
583 + auto local_split_kv = params.split_kv;
584 + if (params.mainloop.ptr_seq != nullptr) {
585 + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)];
586 + if (params.ptr_split_kv != nullptr) {
587 + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)];
588 + }
589 + }
590 + if (local_split_kv <= get<3>(blk_coord))
591 + continue;
592 + load_page_table(
593 + blk_coord,
594 + problem_shape,
595 + params.mainloop,
596 + shared_storage.tensors,
597 + pipeline_page_table, pipeline_pt_producer_state,
598 + local_split_kv
599 + );
600 + }
601 + }
602 + else if (role == WarpRole::kLoad) {
603 + if constexpr (kIsCpAsync) {
604 + CUTLASS_PRAGMA_NO_UNROLL
605 + for (; tile_scheduler.is_valid(); ++tile_scheduler) {
606 + auto blk_coord = tile_scheduler.get_block_coord();
607 + auto problem_shape = params.problem_shape;
608 + auto local_split_kv = params.split_kv;
609 + if (params.mainloop.ptr_seq != nullptr) {
610 + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)];
611 + if (params.ptr_split_kv != nullptr) {
612 + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)];
613 + }
614 + }
615 + if (local_split_kv <= get<3>(blk_coord))
616 + continue;
617 + load_cpasync(
618 + blk_coord,
619 + problem_shape,
620 + params.mainloop,
621 + params.mainloop_params,
622 + shared_storage.tensors,
623 + pipeline_load_qk, pipeline_load_qk_producer_state,
624 + local_split_kv,
625 + /* must be shared pipe */
626 + pipeline_page_table, pipeline_pt_consumer_state
627 + );
628 + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive_and_wait();
629 + }
630 + }
631 + else {
632 + if (params.mainloop.ptr_page_table != nullptr) {
633 + CUTLASS_PRAGMA_NO_UNROLL
634 + for (; tile_scheduler.is_valid(); ++tile_scheduler) {
635 + auto blk_coord = tile_scheduler.get_block_coord();
636 + auto problem_shape = params.problem_shape;
637 + auto local_split_kv = params.split_kv;
638 + if (params.mainloop.ptr_seq != nullptr) {
639 + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)];
640 + if (params.ptr_split_kv != nullptr) {
641 + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)];
642 + }
643 + }
644 + if (local_split_kv <= get<3>(blk_coord))
645 + continue;
646 + load_tma</* paged= */ true>(
647 + blk_coord,
648 + problem_shape,
649 + params.mainloop,
650 + params.mainloop_params,
651 + shared_storage.tensors,
652 + pipeline_load_qk, pipeline_load_qk_producer_state,
653 + pipeline_load_qk, pipeline_load_qk_producer_state,
654 + local_split_kv
655 + );
656 + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive_and_wait();
657 + }
658 + }
659 + else {
660 + CUTLASS_PRAGMA_NO_UNROLL
661 + for (; tile_scheduler.is_valid(); ++tile_scheduler) {
662 + auto blk_coord = tile_scheduler.get_block_coord();
663 + auto problem_shape = params.problem_shape;
664 + auto local_split_kv = params.split_kv;
665 + if (params.mainloop.ptr_seq != nullptr) {
666 + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)];
667 + if (params.ptr_split_kv != nullptr) {
668 + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)];
669 + }
670 + }
671 + if (local_split_kv <= get<3>(blk_coord))
672 + continue;
673 + load_tma<false>(
674 + blk_coord,
675 + problem_shape,
676 + params.mainloop,
677 + params.mainloop_params,
678 + shared_storage.tensors,
679 + pipeline_load_qk, pipeline_load_qk_producer_state,
680 + pipeline_load_qk, pipeline_load_qk_producer_state,
681 + local_split_kv
682 + );
683 + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive_and_wait();
684 + }
685 + }
686 + }
687 + }
688 + else if (role == WarpRole::kMma) {
689 + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
690 + __syncwarp();
691 +
692 + if (is_mma_leader_cta) {
693 + CUTLASS_PRAGMA_NO_UNROLL
694 + for (; tile_scheduler.is_valid(); ++tile_scheduler) {
695 + auto blk_coord = tile_scheduler.get_block_coord();
696 + auto problem_shape = params.problem_shape;
697 + auto local_split_kv = params.split_kv;
698 + if (params.mainloop.ptr_seq != nullptr) {
699 + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)];
700 + if (params.ptr_split_kv != nullptr) {
701 + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)];
702 + }
703 + }
704 + if (local_split_kv <= get<3>(blk_coord))
705 + continue;
706 + mma(blk_coord,
707 + problem_shape,
708 + shared_storage.tensors,
709 + pipeline_load_qk, pipeline_load_qk_consumer_state,
710 + pipeline_load_qk, pipeline_load_qk_consumer_state,
711 + pipeline_mma_s, pipeline_mma_s_producer_state,
712 + pipeline_p_mma, pipeline_p_mma_consumer_state,
713 + pipeline_mma_o, pipeline_mma_o_producer_state,
714 + local_split_kv
715 + );
716 + }
717 + }
718 +
719 + //cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, kNamedBarrierTmemDealloc).arrive_and_wait();
720 +
721 + //uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
722 + //tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
723 + }
724 + else if (role == WarpRole::kCompute) {
725 + CUTLASS_PRAGMA_NO_UNROLL
726 + for (; tile_scheduler.is_valid(); ++tile_scheduler) {
727 + auto blk_coord = tile_scheduler.get_block_coord();
728 + auto problem_shape = params.problem_shape;
729 + auto split_kv = params.split_kv;
730 + auto local_split_kv = split_kv;
731 + if (params.mainloop.ptr_seq != nullptr) {
732 + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)];
733 + if (params.ptr_split_kv != nullptr) {
734 + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)];
735 + }
736 + }
737 + if (local_split_kv <= get<3>(blk_coord))
738 + continue;
739 + compute(
740 + blk_coord,
741 + problem_shape,
742 + params.mainloop, // for softmax_scale
743 + params.epilogue,
744 + shared_storage.tensors, // for smem_comm
745 + pipeline_mma_s, pipeline_mma_s_consumer_state,
746 + pipeline_p_mma, pipeline_p_mma_producer_state,
747 + pipeline_mma_o, pipeline_mma_o_consumer_state,
748 + local_split_kv
749 + );
750 + }
751 +
752 + //cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, kNamedBarrierTmemDealloc).arrive();
753 + }
754 +
755 + cute::cluster_sync();
756 + cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, kNamedBarrierTmemDealloc).arrive();
757 + if (role == WarpRole::kMma) {
758 + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
759 + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
760 + }
761 + }
762 +
763 + template<class BlkCoord>
764 + CUTLASS_DEVICE void load_page_table(
765 + BlkCoord const& blk_coord,
766 + ProblemShape const& problem_shape,
767 + MainloopArguments const& mainloop_args,
768 + TensorStorage& shared_tensors,
769 + PipelinePT& pipeline_page_table,
770 + typename PipelinePT::PipelineState& pipeline_pt_producer_state, int const& split_kv) {
771 +
772 + auto [H, K, D, B] = problem_shape;
773 + int batch_coord = get<2>(blk_coord);
774 +
775 + auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table),
776 + make_shape(mainloop_args.page_count, B),
777 + mainloop_args.stride_page_table);
778 + auto mPT = mPT_l(_, batch_coord);
779 +
780 + int k_tile_total = ceil_div(K, TileShapeS{});
781 + int k_tile_per_cta = ceil_div(k_tile_total, split_kv);
782 + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit
783 + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index);
784 + if (k_tile_count == 0) {
785 + return;
786 + }
787 +
788 + auto page_size = Pow2{mainloop_args.page_size};
789 + auto pages_per_tile = Pow2{TileShapeS{} / page_size};
790 + int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarp;
791 +
792 +#if 1
793 + for (; k_tile_count > 0; ++k_index, --k_tile_count) {
794 + pipeline_page_table.producer_acquire(pipeline_pt_producer_state);
795 +
796 + // assume a single warp
797 +
798 + CUTLASS_PRAGMA_UNROLL
799 + for (int i = 0; i < TileShapeS{}; i += cutlass::NumThreadsPerWarp) {
800 + int idx = i + thread_idx;
801 + bool guard = idx < pages_per_tile;
802 + int smem_idx = pipeline_pt_producer_state.index() * TileShapeS::value + idx;
803 + int pt_idx = pages_per_tile * k_index + idx;
804 +
805 + cutlass::arch::cp_async_zfill<sizeof(int), cutlass::arch::CacheOperation::Always>(
806 + &shared_tensors.smem_page_table[smem_idx], &mPT(pt_idx), guard
807 + );
808 + }
809 +
810 + pipeline_page_table.producer_commit(pipeline_pt_producer_state, cutlass::arch::cpasync_barrier_arrive);
811 + ++pipeline_pt_producer_state;
812 + }
813 +#endif
814 + }
815 +
816 +
817 + struct Gather {
818 + int& page_table_stage;
819 + Pow2 pages_per_tile;
820 + const int * __restrict__ smem_page_table;
821 +
822 + CUTLASS_DEVICE int operator()(int idx) const {
823 + return smem_page_table[page_table_stage * TileShapeS::value + idx % pages_per_tile];
824 + }
825 +
826 + CUTLASS_DEVICE friend void print(Gather const&) {
827 + printf("<gather>");
828 + }
829 +
830 + };
831 +
832 +
833 + template<class BlkCoord>
834 + CUTLASS_DEVICE void load_cpasync(
835 + BlkCoord const& blk_coord,
836 + ProblemShape const& problem_shape,
837 + MainloopArguments const& mainloop_args,
838 + MainloopParams const& mainloop_params,
839 + TensorStorage& shared_tensors,
840 + PipelineLoadQK& pipeline_load,
841 + typename PipelineLoadQK::PipelineState& pipeline_load_producer_state,
842 + int const& split_kv,
843 + PipelinePT& pipeline_page_table,
844 + typename PipelinePT::PipelineState& pipeline_pt_consumer_state) {
845 +
846 + auto [H, K, D, B] = problem_shape;
847 + auto [D_latent, D_rope] = D;
848 +
849 + using X = Underscore;
850 +
851 + int k_tile_total = ceil_div(K, TileShapeS{});
852 + int k_tile_per_cta = ceil_div(k_tile_total, split_kv);
853 + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit
854 + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index);
855 + if (k_tile_count == 0) {
856 + return;
857 + }
858 +
859 + // partition all tensors
860 + auto mQL = make_tensor(make_gmem_ptr(mainloop_args.ptr_q_latent), make_shape(H, D_latent, B), mainloop_args.stride_q_latent);
861 + auto mQR = make_tensor(make_gmem_ptr(mainloop_args.ptr_q_rope), make_shape(H, D_rope, B), mainloop_args.stride_q_rope);
862 +
863 + int paged_B = mainloop_args.page_count;
864 + auto paged_K = Pow2{mainloop_args.page_size};
865 + auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), make_shape(paged_B, B), mainloop_args.stride_page_table);
866 +
867 + int batch_coord = get<2>(blk_coord);
868 + auto mPT = mPT_l(_, batch_coord);
869 +
870 + auto gQL = local_tile(mQL, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{});
871 + auto gQR = local_tile(mQR, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{});
872 +
873 + ThrMMA cta_mma_qk = TiledMmaQK{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{}));
874 + ThrMMA cta_mma_pv = TiledMmaPV{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{}));
875 +
876 + auto tSgQL = cta_mma_qk.partition_A(gQL);
877 + auto tSgQR = cta_mma_qk.partition_A(gQR);
878 +
879 + Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{});
880 + Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{});
881 + Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{});
882 +
883 + auto make_copy_for = [](auto sT) {
884 + auto rT_a = sT.layout()(_, _, _, _0{});
885 + auto rT = make_ordered_layout(shape(rT_a), stride(rT_a));
886 + auto threads = Int<kNumLoadWarps * cutlass::NumThreadsPerWarp>{};
887 + auto values = Int<sizeof(uint128_t) / sizeof(Element)>{};
888 + return make_cotiled_copy(
889 + Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, Element>{},
890 + make_ordered_layout(
891 + make_shape(threads, values),
892 + make_stride(_1{}, _0{})),
893 + rT);
894 + };
895 +
896 + // like cute::copy, but makes sure we do all page table lookups first
897 + auto copy_split = [](auto atom, auto src, auto dst) {
898 + auto src_v = group_modes<1, rank_v<decltype(src)>>(src);
899 + auto dst_v = group_modes<1, rank_v<decltype(dst)>>(dst);
900 +
901 + auto src_v_ptrs = make_tensor<Element*>(size<1>(src_v));
902 + for (int i = 0; i < size<1>(src_v); i++) {
903 + src_v_ptrs(i) = &src_v(_0{}, i);
904 + }
905 +
906 +
907 + for (int i = 0; i < size<1>(src_v); i++) {
908 + auto src_v_i = make_tensor(
909 + make_gmem_ptr(src_v_ptrs(i)),
910 + make_shape(shape<0>(src_v)),
911 + make_stride(make_stride(_1{}, _0{}))
912 + );
913 + atom.call(src_v_i, dst_v(_, i));
914 + }
915 + };
916 +
917 + auto tiled_copy_q = make_copy_for(sQ);
918 + auto tiled_copy_kc = make_copy_for(sKC);
919 + auto tiled_copy_vc = make_copy_for(sVC);
920 +
921 + auto thr_copy_q = tiled_copy_q.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp));
922 + auto thr_copy_kc = tiled_copy_kc.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp));
923 + auto thr_copy_vc = tiled_copy_vc.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp));
924 +
925 + auto tQsQ = thr_copy_q.partition_D(sQ);
926 + auto tQgQL = thr_copy_q.partition_S(tSgQL);
927 + auto tQgQR = thr_copy_q.partition_S(tSgQR);
928 +
929 + auto tKCsKC = thr_copy_kc.partition_D(sKC);
930 + auto tVCsVC = thr_copy_vc.partition_D(sVC);
931 +
932 + auto pipeline_pt_release_state = pipeline_pt_consumer_state;
933 +
934 + int page_table_stage = -1;
935 + Pow2 pages_per_tile{TileShapeS{} / paged_K};
936 + const int * __restrict__ smem_page_table = shared_tensors.smem_page_table.begin();
937 + Gather gather{page_table_stage, pages_per_tile, smem_page_table};
938 +
939 + auto mCL = make_tensor(
940 + make_gmem_ptr(mainloop_args.ptr_c_latent),
941 + ComposedLayout{
942 + make_layout(
943 + make_shape(make_shape(paged_K, paged_B), _1{}),
944 + make_stride(make_stride(get<0>(mainloop_args.stride_c_latent), example::CustomStride(gather, get<2>(mainloop_args.stride_c_latent))), get<1>(mainloop_args.stride_c_latent))),
945 + make_coord(_0{}, _0{}),
946 + make_identity_layout(make_shape(paged_K * paged_B, D_latent))});
947 +
948 + auto mKR = make_tensor(
949 + make_gmem_ptr(mainloop_args.ptr_k_rope),
950 + ComposedLayout{
951 + make_layout(
952 + make_shape(make_shape(paged_K, paged_B), _1{}),
953 + make_stride(make_stride(get<0>(mainloop_args.stride_k_rope), example::CustomStride(gather, get<2>(mainloop_args.stride_k_rope))), get<1>(mainloop_args.stride_k_rope))),
954 + make_coord(_0{}, _0{}),
955 + make_identity_layout(make_shape(paged_K * paged_B, D_latent))});
956 +
957 + auto mCLT = make_tensor(
958 + make_gmem_ptr(mainloop_args.ptr_c_latent),
959 + ComposedLayout{
960 + make_layout(
961 + make_shape(_1{}, make_shape(paged_K, paged_B)),
962 + make_stride(get<1>(mainloop_args.stride_c_latent), make_stride(get<0>(mainloop_args.stride_c_latent), example::CustomStride(gather, get<2>(mainloop_args.stride_c_latent))))),
963 + make_coord(_0{}, _0{}),
964 + make_identity_layout(make_shape(D_latent, paged_K * paged_B))});
965 +
966 + auto gCL = local_tile(mCL, TileShapeQK{}, make_coord(_,_,_), Step<X, _1, _1>{});
967 + auto gKR = local_tile(mKR, TileShapeQK{}, make_coord(_,_,_), Step<X, _1, _1>{});
968 + auto gCLT = local_tile(mCLT, TileShapePV{}, make_coord(_,_,_), Step<X, _1, _1>{});
969 +
970 + auto tSgCL = cta_mma_qk.partition_B(gCL);
971 + auto tSgKR = cta_mma_qk.partition_B(gKR);
972 + auto tOgCLT = cta_mma_pv.partition_B(gCLT);
973 +
974 + auto tKCgCL = thr_copy_kc.partition_S(tSgCL);
975 + auto tKCgKR = thr_copy_kc.partition_S(tSgKR);
976 + auto tVCgCLT = thr_copy_vc.partition_S(tOgCLT);
977 +
978 + // latent is first in memory, so let's load it first always
979 + // startup: alternate Q and K, set tx count appropriately, for k_idx = 0
980 + auto& pipeline_acquire_state = pipeline_load_producer_state;
981 + auto pipeline_commit_state = pipeline_acquire_state;
982 + int pipeline_offset = 0;
983 +
984 + for (int i = 0; i < StagesPV; i++) {
985 + cutlass::arch::cp_async_fence();
986 + }
987 +
988 + auto load_stage = [&](auto fn) {
989 + pipeline_load.producer_acquire(pipeline_acquire_state);
990 + fn(pipeline_acquire_state.index());
991 + cutlass::arch::cp_async_fence();
992 +
993 + ++pipeline_acquire_state;
994 + ++pipeline_offset;
995 +
996 + if (pipeline_offset == StagesPV - 1) {
997 + cutlass::arch::cp_async_wait<StagesPV - 1>();
998 + pipeline_load.producer_commit(pipeline_commit_state);
999 + ++pipeline_commit_state;
1000 + --pipeline_offset;
1001 + }
1002 + };
1003 +
1004 + pipeline_page_table.consumer_wait(pipeline_pt_consumer_state);
1005 + page_table_stage = pipeline_pt_consumer_state.index();
1006 + ++pipeline_pt_consumer_state;
1007 +
1008 + // each Q/K tile consists of rope and latent
1009 + for (int i = 0; i < IterationsQKLatent; i++) {
1010 + load_stage([&](int index) {
1011 + cute::copy(tiled_copy_q, tQgQL(_, _, _, _, _0{}, i, batch_coord), tQsQ(_, _, _, _, i));
1012 + copy_split(tiled_copy_kc, tKCgCL(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index));
1013 + });
1014 + }
1015 +
1016 + for (int i = 0; i < IterationsQKRope; i++) {
1017 + load_stage([&](int index) {
1018 + cute::copy(tiled_copy_q, tQgQR(_, _, _, _, _0{}, i, batch_coord), tQsQ(_, _, _, _, IterationsQKLatent + i));
1019 + copy_split(tiled_copy_kc, tKCgKR(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index));
1020 + });
1021 + }
1022 +
1023 + k_index += 1;
1024 + k_tile_count -= 1;
1025 +
1026 + // assume k_tile_count >= 1
1027 + // perform K+Q load here
1028 + CUTLASS_PRAGMA_NO_UNROLL
1029 + while (k_tile_count > 0) {
1030 +
1031 + pipeline_page_table.consumer_wait(pipeline_pt_consumer_state);
1032 + page_table_stage = pipeline_pt_consumer_state.index();
1033 + ++pipeline_pt_consumer_state;
1034 +
1035 + for (int i = 0; i < IterationsQKLatent; i++) {
1036 + load_stage([&](int index) {
1037 + copy_split(tiled_copy_kc, tKCgCL(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index));
1038 + });
1039 + }
1040 +
1041 + for (int i = 0; i < IterationsQKRope; i++) {
1042 + load_stage([&](int index) {
1043 + copy_split(tiled_copy_kc, tKCgKR(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index));
1044 + });
1045 + }
1046 +
1047 + page_table_stage = pipeline_pt_release_state.index();
1048 +
1049 + for (int i = 0; i < IterationsPV_K; i++) {
1050 + for (int j = 0; j < IterationsPV_N; j++) {
1051 + load_stage([&](int index) {
1052 + copy_split(tiled_copy_vc, tVCgCLT(_, _, _, _, j, IterationsPV_K * (k_index - 1) + i), tVCsVC(_, _, _, _, index));
1053 + });
1054 + }
1055 + }
1056 +
1057 + pipeline_page_table.consumer_release(pipeline_pt_release_state);
1058 + ++pipeline_pt_release_state;
1059 +
1060 + k_index += 1;
1061 + k_tile_count -= 1;
1062 + }
1063 +
1064 + page_table_stage = pipeline_pt_release_state.index();
1065 +
1066 + for (int i = 0; i < IterationsPV_K; i++) {
1067 + for (int j = 0; j < IterationsPV_N; j++) {
1068 + load_stage([&](int index) {
1069 + copy_split(tiled_copy_vc, tVCgCLT(_, _, _, _, j, IterationsPV_K * (k_index - 1) + i), tVCsVC(_, _, _, _, index));
1070 + });
1071 + }
1072 + }
1073 +
1074 + pipeline_page_table.consumer_release(pipeline_pt_release_state);
1075 + ++pipeline_pt_release_state;
1076 +
1077 + while (pipeline_offset > 0) {
1078 + cutlass::arch::cp_async_fence();
1079 +
1080 + cutlass::arch::cp_async_wait<StagesPV - 1>();
1081 + pipeline_load.producer_commit(pipeline_commit_state);
1082 + ++pipeline_commit_state;
1083 + --pipeline_offset;
1084 + }
1085 +
1086 + cutlass::arch::cp_async_wait<0>();
1087 +
1088 + }
1089 +
1090 +
1091 + template<bool kIsPaged = false, class BlkCoord>
1092 + CUTLASS_DEVICE void load_tma(
1093 + BlkCoord const& blk_coord,
1094 + ProblemShape const& problem_shape,
1095 + MainloopArguments const& mainloop_args,
1096 + MainloopParams const& mainloop_params,
1097 + TensorStorage& shared_tensors,
1098 + PipelineLoadQK& pipeline_load_qk,
1099 + typename PipelineLoadQK::PipelineState& pipeline_load_qk_producer_state,
1100 + PipelineLoadPV& pipeline_load_pv,
1101 + typename PipelineLoadPV::PipelineState& pipeline_load_pv_producer_state,
1102 + int const& split_kv) {
1103 +
1104 + auto [H, K, D, B] = problem_shape;
1105 + auto [D_latent, D_rope] = D;
1106 +
1107 + int k_tile_total = ceil_div(K, TileShapeS{});
1108 + int k_tile_per_cta = ceil_div(k_tile_total, split_kv);
1109 + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit
1110 + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index);
1111 + if (k_tile_count == 0) {
1112 + return;
1113 + }
1114 +
1115 + using X = Underscore;
1116 +
1117 + // partition all tensors
1118 + auto mQL = mainloop_params.tma_load_q_latent.get_tma_tensor(make_shape(H, D_latent, B));
1119 + auto mQR = mainloop_params.tma_load_q_rope.get_tma_tensor(make_shape(H, D_rope, B));
1120 +
1121 + int paged_B = B;
1122 + int paged_K = K;
1123 + if constexpr (kIsPaged) {
1124 + paged_B = mainloop_args.page_count;
1125 + paged_K = mainloop_args.page_size;
1126 + }
1127 + auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), make_shape(paged_B, B), mainloop_args.stride_page_table);
1128 +
1129 + auto mCL = mainloop_params.tma_load_c_latent.get_tma_tensor(make_shape(paged_K, D_latent, paged_B));
1130 + auto mKR = mainloop_params.tma_load_k_rope.get_tma_tensor(make_shape(paged_K, D_rope, paged_B));
1131 +
1132 + auto mCLT = mainloop_params.tma_load_c_latent_transpose.get_tma_tensor(make_shape(D_latent, paged_K, paged_B));
1133 +
1134 + auto gQL = local_tile(mQL, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{});
1135 + auto gQR = local_tile(mQR, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{});
1136 +
1137 + auto gCL = local_tile(mCL, TileShapeQK{}, make_coord(_,_,_), Step<X, _1, _1>{});
1138 + auto gKR = local_tile(mKR, TileShapeQK{}, make_coord(_,_,_), Step<X, _1, _1>{});
1139 + auto gCLT = local_tile(mCLT, TileShapePV{}, make_coord(_,_,_), Step<X, _1, _1>{});
1140 +
1141 + ThrMMA cta_mma_qk = TiledMmaQK{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{}));
1142 + ThrMMA cta_mma_pv = TiledMmaPV{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{}));
1143 +
1144 + auto tSgQL = cta_mma_qk.partition_A(gQL);
1145 + auto tSgQR = cta_mma_qk.partition_A(gQR);
1146 +
1147 + auto tSgCL = cta_mma_qk.partition_B(gCL);
1148 + auto tSgKR = cta_mma_qk.partition_B(gKR);
1149 +
1150 + auto tOgCLT = cta_mma_pv.partition_B(gCLT);
1151 +
1152 + Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{});
1153 + Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{});
1154 + Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{});
1155 +
1156 + auto [tQLgQL_mkl, tQsQ] = tma_partition(
1157 + mainloop_params.tma_load_q_latent, _0{}, make_layout(_1{}),
1158 + group_modes<0,3>(sQ), group_modes<0,3>(tSgQL));
1159 +
1160 + auto [tQRgQR_mkl, tQsQ_ignore] = tma_partition(
1161 + mainloop_params.tma_load_q_rope, _0{}, make_layout(_1{}),
1162 + group_modes<0,3>(sQ), group_modes<0,3>(tSgQR));
1163 +
1164 + auto [tCLgCL_nkl, tKCsKC] = tma_partition(
1165 + mainloop_params.tma_load_c_latent, _0{}, make_layout(_1{}),
1166 + group_modes<0,3>(sKC), group_modes<0,3>(tSgCL));
1167 +
1168 + auto [tKRgKR_nkl, tKCsKC_ignore] = tma_partition(
1169 + mainloop_params.tma_load_k_rope, _0{}, make_layout(_1{}),
1170 + group_modes<0,3>(sKC), group_modes<0,3>(tSgKR));
1171 +
1172 + auto [tCLTgCLT_nkl, tVCsVC] = tma_partition(
1173 + mainloop_params.tma_load_c_latent_transpose, _0{}, make_layout(_1{}),
1174 + group_modes<0,3>(sVC), group_modes<0,3>(tOgCLT));
1175 +
1176 + uint16_t mcast_mask = 0;
1177 +
1178 + int batch_coord = get<2>(blk_coord);
1179 + Tensor tQLgQL = tQLgQL_mkl(_, _, _, batch_coord);
1180 + Tensor tQRgQR = tQRgQR_mkl(_, _, _, batch_coord);
1181 +
1182 + auto mPT = mPT_l(_, batch_coord);
1183 +
1184 + Tensor tCLgCL = tCLgCL_nkl(_, _, _, _);
1185 + Tensor tKRgKR = tKRgKR_nkl(_, _, _, _);
1186 +
1187 + // careful: stage and k are swapped here!
1188 + Tensor tCLTgCLT = tCLTgCLT_nkl(_, _, _, _);
1189 +
1190 + // latent is first in memory, so let's load it first always
1191 + // startup: alternate Q and K, set tx count appropriately, for k_idx = 0
1192 +
1193 + // each Q/K tile consists of rope and latent
1194 + for (int i = 0; i < IterationsQKLatent; i++) {
1195 + pipeline_load_qk.producer_expect_transaction(pipeline_load_qk_producer_state, kTransactionsBytesLoadExtraQ);
1196 + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state);
1197 + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state);
1198 +
1199 + if (cute::elect_one_sync()) {
1200 + // expect the extra bytes
1201 + // load_qk ql
1202 + cute::copy(mainloop_params.tma_load_q_latent.with(*tma_barrier, mcast_mask), tQLgQL(_, _0{}, i), tQsQ(_, i));
1203 + // load_qk cl
1204 + if constexpr (kIsPaged) {
1205 + cute::copy(
1206 + mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask),
1207 + tCLgCL(_, _0{}, i, mPT(k_index)),
1208 + tKCsKC(_, pipeline_load_qk_producer_state.index())
1209 + );
1210 + }
1211 + else {
1212 + cute::copy(
1213 + mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask),
1214 + tCLgCL(_, k_index, i, batch_coord),
1215 + tKCsKC(_, pipeline_load_qk_producer_state.index()));
1216 + }
1217 + }
1218 + ++pipeline_load_qk_producer_state;
1219 + }
1220 +
1221 + for (int i = 0; i < IterationsQKRope; i++) {
1222 + pipeline_load_qk.producer_expect_transaction(pipeline_load_qk_producer_state, kTransactionsBytesLoadExtraQ);
1223 + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state);
1224 + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state);
1225 +
1226 + if (cute::elect_one_sync()) {
1227 + // expect the extra bytes
1228 + // load_qk ql
1229 + cute::copy(mainloop_params.tma_load_q_rope.with(*tma_barrier, mcast_mask), tQRgQR(_, _0{}, i), tQsQ(_, i + IterationsQKLatent));
1230 + // load_qk cl
1231 + if constexpr (kIsPaged) {
1232 + cute::copy(
1233 + mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask),
1234 + tKRgKR(_, _0{}, i, mPT(k_index)),
1235 + tKCsKC(_, pipeline_load_qk_producer_state.index())
1236 + );
1237 + }
1238 + else {
1239 + cute::copy(
1240 + mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask),
1241 + tKRgKR(_, k_index, i, batch_coord),
1242 + tKCsKC(_, pipeline_load_qk_producer_state.index()));
1243 + }
1244 + }
1245 + ++pipeline_load_qk_producer_state;
1246 + }
1247 +
1248 + k_index += 1;
1249 + k_tile_count -= 1;
1250 +
1251 + // assume k_tile_count >= 1
1252 + // perform K+Q load here
1253 + CUTLASS_PRAGMA_NO_UNROLL
1254 + while (k_tile_count > 0) {
1255 +
1256 + // perform K load
1257 + for (int i = 0; i < IterationsQKLatent; i++) {
1258 + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state);
1259 + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state);
1260 +
1261 + if (cute::elect_one_sync()) {
1262 + // load_qk cl
1263 + if constexpr (kIsPaged) {
1264 + cute::copy(
1265 + mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask),
1266 + tCLgCL(_, _0{}, i, mPT(k_index)),
1267 + tKCsKC(_, pipeline_load_qk_producer_state.index())
1268 + );
1269 + }
1270 + else {
1271 + cute::copy(
1272 + mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask),
1273 + tCLgCL(_, k_index, i, batch_coord),
1274 + tKCsKC(_, pipeline_load_qk_producer_state.index()));
1275 + }
1276 + }
1277 + ++pipeline_load_qk_producer_state;
1278 + }
1279 +
1280 + for (int i = 0; i < IterationsQKRope; i++) {
1281 + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state);
1282 + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state);
1283 +
1284 + if (cute::elect_one_sync()) {
1285 + // load_qk cl
1286 + if constexpr (kIsPaged) {
1287 + cute::copy(
1288 + mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask),
1289 + tKRgKR(_, _0{}, i, mPT(k_index)),
1290 + tKCsKC(_, pipeline_load_qk_producer_state.index())
1291 + );
1292 + }
1293 + else {
1294 + cute::copy(
1295 + mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask),
1296 + tKRgKR(_, k_index, i, batch_coord),
1297 + tKCsKC(_, pipeline_load_qk_producer_state.index()));
1298 + }
1299 + }
1300 + ++pipeline_load_qk_producer_state;
1301 + }
1302 +
1303 + // prefetch next K load to keep busy while we transpose-load from cache
1304 + const int kPrefetchDistance = 1;
1305 + for (int i = 0; i < IterationsQKLatent; i++) {
1306 + if (cute::elect_one_sync()) {
1307 + if constexpr (kIsPaged) {
1308 + if (k_tile_count > kPrefetchDistance) {
1309 + cute::prefetch(
1310 + mainloop_params.tma_load_c_latent,
1311 + tCLgCL(_, _0{}, i, mPT(k_index + kPrefetchDistance))
1312 + );
1313 + }
1314 + }
1315 + else {
1316 + cute::prefetch(
1317 + mainloop_params.tma_load_c_latent,
1318 + tCLgCL(_, k_index + kPrefetchDistance, i, batch_coord)
1319 + );
1320 + }
1321 + }
1322 + }
1323 +
1324 + for (int i = 0; i < IterationsQKRope; i++) {
1325 + if (cute::elect_one_sync()) {
1326 + if constexpr (kIsPaged) {
1327 + if (k_tile_count > kPrefetchDistance) {
1328 + cute::prefetch(
1329 + mainloop_params.tma_load_k_rope,
1330 + tKRgKR(_, _0{}, i, mPT(k_index + kPrefetchDistance))
1331 + );
1332 + }
1333 + }
1334 + else {
1335 + cute::prefetch(
1336 + mainloop_params.tma_load_k_rope,
1337 + tKRgKR(_, k_index + kPrefetchDistance, i, batch_coord)
1338 + );
1339 + }
1340 + }
1341 + }
1342 +
1343 + // perform V load (k_idx - 1)
1344 +
1345 + for (int i = 0; i < IterationsPV_K; i++) {
1346 + for (int j = 0; j < IterationsPV_N; j++) {
1347 + pipeline_load_pv.producer_acquire(pipeline_load_pv_producer_state);
1348 + auto tma_barrier = pipeline_load_pv.producer_get_barrier(pipeline_load_pv_producer_state);
1349 +
1350 + if (cute::elect_one_sync()) {
1351 + // load_pv cl
1352 + // note the transpose in indices!
1353 + // note we are off-by-one on k_index
1354 + if constexpr (kIsPaged) {
1355 + cute::copy(
1356 + mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST),
1357 + tCLTgCLT(_, j, i, mPT(k_index - 1)),
1358 + tVCsVC(_, pipeline_load_pv_producer_state.index())
1359 + );
1360 + }
1361 + else {
1362 + cute::copy(
1363 + mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST),
1364 + tCLTgCLT(_, j, IterationsPV_K * (k_index - 1) + i, batch_coord),
1365 + tVCsVC(_, pipeline_load_pv_producer_state.index())
1366 + );
1367 + }
1368 + }
1369 + ++pipeline_load_pv_producer_state;
1370 + }
1371 + }
1372 +
1373 + k_index += 1;
1374 + k_tile_count -= 1;
1375 + }
1376 +
1377 + for (int i = 0; i < IterationsPV_K; i++) {
1378 + for (int j = 0; j < IterationsPV_N; j++) {
1379 + pipeline_load_pv.producer_acquire(pipeline_load_pv_producer_state);
1380 + auto tma_barrier = pipeline_load_pv.producer_get_barrier(pipeline_load_pv_producer_state);
1381 +
1382 + if (cute::elect_one_sync()) {
1383 + // load_pv cl
1384 + // note the transpose in indices
1385 + // note we are off-by-one on k_index
1386 +
1387 + if constexpr (kIsPaged) {
1388 + cute::copy(
1389 + mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST),
1390 + tCLTgCLT(_, j, i, mPT(k_index - 1)),
1391 + tVCsVC(_, pipeline_load_pv_producer_state.index())
1392 + );
1393 + }
1394 + else {
1395 + cute::copy(
1396 + mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST),
1397 + tCLTgCLT(_, j, IterationsPV_K * (k_index - 1) + i, batch_coord),
1398 + tVCsVC(_, pipeline_load_pv_producer_state.index())
1399 + );
1400 + }
1401 + }
1402 + ++pipeline_load_pv_producer_state;
1403 + }
1404 + }
1405 + }
1406 +
1407 + template<class BlkCoord>
1408 + CUTLASS_DEVICE void mma(
1409 + BlkCoord const& blk_coord,
1410 + ProblemShape const& problem_shape,
1411 + TensorStorage& shared_tensors,
1412 + PipelineLoadQK& pipeline_load_qk,
1413 + typename PipelineLoadQK::PipelineState& pipeline_load_qk_consumer_state,
1414 + PipelineLoadPV& pipeline_load_pv,
1415 + typename PipelineLoadPV::PipelineState& pipeline_load_pv_consumer_state,
1416 + PipelineS& pipeline_mma_s,
1417 + typename PipelineS::PipelineState& pipeline_mma_s_producer_state,
1418 + PipelineP& pipeline_p_mma,
1419 + typename PipelineP::PipelineState& pipeline_p_mma_consumer_state,
1420 + PipelineO& pipeline_mma_o,
1421 + typename PipelineO::PipelineState& pipeline_mma_o_producer_state,
1422 + int const& split_kv) {
1423 +
1424 + auto [H, K, D, B] = problem_shape;
1425 +
1426 + int k_tile_total = ceil_div(K, TileShapeS{});
1427 + int k_tile_per_cta = ceil_div(k_tile_total, split_kv);
1428 + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit
1429 + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index);
1430 + if (k_tile_count == 0) {
1431 + return;
1432 + }
1433 +
1434 + // mma init
1435 + Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{});
1436 + Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{});
1437 + Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{});
1438 + Tensor sP = make_tensor(make_smem_ptr((Element*) shared_tensors.smem_p.begin()), SmemLayoutP{});
1439 +
1440 + Tensor tSrQ = TiledMmaQK::make_fragment_A(sQ);
1441 + Tensor tSrKC = TiledMmaQK::make_fragment_B(sKC);
1442 + Tensor tOrP = TiledMmaPV::make_fragment_A(sP);
1443 + Tensor tOrVC = TiledMmaPV::make_fragment_B(sVC);
1444 +
1445 + TiledMmaQK tiled_mma_qk;
1446 + TiledMmaPV tiled_mma_pv;
1447 +
1448 + Tensor tStS = partition_fragment_C(tiled_mma_qk, select<0,1>(TileShapeQK{}));
1449 + Tensor tItI = partition_fragment_C(tiled_mma_pv, select<0,1>(TileShapePV{}));
1450 +
1451 + tiled_mma_pv.accumulate_ = UMMA::ScaleOut::Zero;
1452 +
1453 + pipeline_mma_s.producer_acquire(pipeline_mma_s_producer_state);
1454 +
1455 + // Mma S0 S1 O0 S2 O1 ... Sn On-1 On
1456 + // S0 ownership -- ----- -- --
1457 + // S1 ownership -- ----- ----
1458 + // O ownership -- -- ---- --
1459 +
1460 + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::Zero;
1461 + for (int i = 0; i < IterationsQK; i++) {
1462 + pipeline_load_qk.consumer_wait(pipeline_load_qk_consumer_state);
1463 + int read_stage = pipeline_load_qk_consumer_state.index();
1464 +
1465 + tStS.data() = uint32_t(pipeline_mma_s_producer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1);
1466 +
1467 + CUTLASS_PRAGMA_UNROLL
1468 + for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) {
1469 + cute::gemm(tiled_mma_qk,
1470 + tSrQ(_,_,k_block,i),
1471 + tSrKC(_,_,k_block,read_stage),
1472 + tStS);
1473 + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::One;
1474 + }
1475 +
1476 + pipeline_load_qk.consumer_release(pipeline_load_qk_consumer_state);
1477 + ++pipeline_load_qk_consumer_state;
1478 + }
1479 +
1480 + pipeline_mma_s.producer_commit(pipeline_mma_s_producer_state);
1481 + ++pipeline_mma_s_producer_state;
1482 +
1483 + k_tile_count -= 1;
1484 +
1485 + CUTLASS_PRAGMA_NO_UNROLL
1486 + while (k_tile_count > 0) {
1487 +
1488 + pipeline_mma_s.producer_acquire(pipeline_mma_s_producer_state);
1489 + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::Zero;
1490 + for (int i = 0; i < IterationsQK; i++) {
1491 + pipeline_load_qk.consumer_wait(pipeline_load_qk_consumer_state);
1492 + int read_stage = pipeline_load_qk_consumer_state.index();
1493 +
1494 + tStS.data() = uint32_t(pipeline_mma_s_producer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1);
1495 +
1496 + CUTLASS_PRAGMA_UNROLL
1497 + for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) {
1498 + cute::gemm(tiled_mma_qk,
1499 + tSrQ(_,_,k_block,i),
1500 + tSrKC(_,_,k_block,read_stage),
1501 + tStS);
1502 + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::One;
1503 + }
1504 +
1505 + pipeline_load_qk.consumer_release(pipeline_load_qk_consumer_state);
1506 + ++pipeline_load_qk_consumer_state;
1507 + }
1508 +
1509 + pipeline_mma_s.producer_commit(pipeline_mma_s_producer_state);
1510 + ++pipeline_mma_s_producer_state;
1511 +
1512 + pipeline_mma_o.producer_acquire(pipeline_mma_o_producer_state);
1513 + pipeline_p_mma.consumer_wait(pipeline_p_mma_consumer_state);
1514 +
1515 + for (int i = 0; i < IterationsPV_K; i++) {
1516 + auto acc_flag = tiled_mma_pv.accumulate_;
1517 + for (int j = 0; j < IterationsPV_N; j++) {
1518 + pipeline_load_pv.consumer_wait(pipeline_load_pv_consumer_state);
1519 +
1520 + int read_stage = pipeline_load_pv_consumer_state.index();
1521 +
1522 + tItI.data() = uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO);
1523 + tiled_mma_pv.accumulate_ = acc_flag;
1524 +
1525 + CUTLASS_PRAGMA_UNROLL
1526 + for (int k_block = 0; k_block < size<2>(tOrP); ++k_block) {
1527 + cute::gemm(tiled_mma_pv,
1528 + tOrP(_,_,k_block, make_coord(i, pipeline_p_mma_consumer_state.index())),
1529 + tOrVC(_,_,k_block,read_stage),
1530 + tItI);
1531 + tiled_mma_pv.accumulate_ = UMMA::ScaleOut::One;
1532 + }
1533 +
1534 + pipeline_load_pv.consumer_release(pipeline_load_pv_consumer_state);
1535 + ++pipeline_load_pv_consumer_state;
1536 + }
1537 + }
1538 +
1539 + pipeline_p_mma.consumer_release(pipeline_p_mma_consumer_state);
1540 + ++pipeline_p_mma_consumer_state;
1541 + pipeline_mma_o.producer_commit(pipeline_mma_o_producer_state);
1542 + ++pipeline_mma_o_producer_state;
1543 +
1544 + --k_tile_count;
1545 + }
1546 +
1547 + pipeline_mma_o.producer_acquire(pipeline_mma_o_producer_state);
1548 + pipeline_p_mma.consumer_wait(pipeline_p_mma_consumer_state);
1549 +
1550 + for (int i = 0; i < IterationsPV_K; i++) {
1551 + auto acc_flag = tiled_mma_pv.accumulate_;
1552 + for (int j = 0; j < IterationsPV_N; j++) {
1553 + pipeline_load_pv.consumer_wait(pipeline_load_pv_consumer_state);
1554 +
1555 + int read_stage = pipeline_load_pv_consumer_state.index();
1556 +
1557 + tItI.data() = uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO);
1558 + tiled_mma_pv.accumulate_ = acc_flag;
1559 +
1560 + CUTLASS_PRAGMA_UNROLL
1561 + for (int k_block = 0; k_block < size<2>(tOrP); ++k_block) {
1562 + cute::gemm(tiled_mma_pv,
1563 + tOrP(_,_,k_block, make_coord(i, pipeline_p_mma_consumer_state.index())),
1564 + tOrVC(_,_,k_block,read_stage),
1565 + tItI);
1566 + tiled_mma_pv.accumulate_ = UMMA::ScaleOut::One;
1567 + }
1568 +
1569 + pipeline_load_pv.consumer_release(pipeline_load_pv_consumer_state);
1570 + ++pipeline_load_pv_consumer_state;
1571 + }
1572 + }
1573 +
1574 + pipeline_p_mma.consumer_release(pipeline_p_mma_consumer_state);
1575 + ++pipeline_p_mma_consumer_state;
1576 + pipeline_mma_o.producer_commit(pipeline_mma_o_producer_state);
1577 + ++pipeline_mma_o_producer_state;
1578 + }
1579 +
1580 +
1581 + template<class IsLastTile>
1582 + CUTLASS_DEVICE void softmax(
1583 + IsLastTile const& is_last_tile,
1584 + ElementAcc& row_max,
1585 + ElementAcc& row_sum,
1586 + ElementAcc& correction_factor,
1587 + ProblemShape const& problem_shape,
1588 + MainloopArguments const& mainloop_args,
1589 + TensorStorage& shared_tensors,
1590 + int k_index,
1591 + uint32_t tmem_s,
1592 + int smem_p_index) {
1593 +
1594 + auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{};
1595 +
1596 + TiledMmaQK tiled_mma_qk;
1597 +
1598 + Tensor tStS = partition_fragment_C(tiled_mma_qk, select<0,1>(TileShapeQK{}));
1599 + tStS.data() = tmem_s;
1600 +
1601 + CUTE_STATIC_ASSERT_V(shape<1>(tStS) == _1{});
1602 + CUTE_STATIC_ASSERT_V(shape<2>(tStS) == _1{});
1603 + Tensor tAcc = tStS(make_coord(_,_),_0{},_0{});
1604 +
1605 + Tensor cS = make_identity_tensor(take<0,2>(CtaShapeQK{}));
1606 +
1607 + auto tiled_t2r = make_tmem_copy(load_op, tAcc);
1608 + auto thread_idx = threadIdx.x % size(tiled_t2r);
1609 +
1610 + auto thread_t2r = tiled_t2r.get_slice(thread_idx);
1611 + Tensor tTR_cS = thread_t2r.partition_D(cS);
1612 + Tensor tTR_rAcc = make_tensor<ElementAcc>(shape(tTR_cS));
1613 +
1614 + Tensor tTR_rS_frag = make_tensor<Element>(shape(tTR_rAcc));
1615 + const int AlignmentS = 4;
1616 + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc);
1617 + Tensor tTR_rAcc_vec = recast<Array<ElementAcc, AlignmentS>>(tTR_rAcc);
1618 + Tensor tTR_rS_vec = recast<Array<Element, AlignmentS>>(tTR_rS_frag);
1619 +
1620 + // load s
1621 + copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
1622 +
1623 + if (is_last_tile) {
1624 + for (int i = 0; i < size(tTR_rAcc); i++) {
1625 + if (get<1>(tTR_cS(i)) + TileShapeS{} * k_index >= get<1>(problem_shape)) {
1626 + tTR_rAcc(i) = -std::numeric_limits<ElementAcc>::infinity();
1627 + }
1628 + }
1629 + }
1630 +
1631 + // max
1632 + ElementAcc row_max_new = row_max;
1633 + CUTLASS_PRAGMA_UNROLL
1634 + for (int i = 0; i < size(tTR_rAcc); i += 1) {
1635 + row_max_new = ::fmax(row_max_new, tTR_rAcc(i));
1636 + }
1637 +
1638 + // for 2x2 dp, reduce here
1639 + if constexpr (kWarpsInN > 1) {
1640 + shared_tensors.smem_exchange[threadIdx.x] = row_max_new;
1641 + cutlass::arch::NamedBarrier(kNumComputeWarps*NumThreadsPerWarp, kNamedBarrierExchange).sync();
1642 + // (64, 2) shape
1643 + int peer_index = (threadIdx.x + 64) % 128;
1644 + row_max_new = cutlass::max(row_max_new, shared_tensors.smem_exchange[peer_index]);
1645 + }
1646 +
1647 +#ifndef B2B
1648 + // find correction factor
1649 + ElementAcc softmax_scale_log2 = mainloop_args.softmax_scale * static_cast<ElementAcc>(M_LOG2E);
1650 + correction_factor = ::exp2f(softmax_scale_log2 * (row_max - row_max_new));
1651 + row_max = row_max_new;
1652 +
1653 + // softmax
1654 + ElementAcc row_max_scale_log2 = row_max * softmax_scale_log2;
1655 + CUTLASS_PRAGMA_UNROLL
1656 + for (int i = 0; i < size(tTR_rAcc); i++) {
1657 + tTR_rAcc(i) = ::exp2f(softmax_scale_log2 * tTR_rAcc(i) - row_max_scale_log2);
1658 + }
1659 +#endif
1660 +
1661 + // quantize
1662 + cutlass::NumericArrayConverter<Element, ElementAcc, AlignmentS> epilogue_op;
1663 +
1664 + CUTLASS_PRAGMA_UNROLL
1665 + for (int i = 0; i < size(tTR_rAcc_vec); i++) {
1666 + tTR_rS_vec(i) = epilogue_op(tTR_rAcc_vec(i));
1667 + }
1668 +
1669 + Tensor sP = make_tensor(make_smem_ptr((Element*) shared_tensors.smem_p.begin()), SmemLayoutP{})(_, _, _, make_coord(_, smem_p_index));
1670 +
1671 + Tensor tOcP = TiledMmaPV{}.get_slice(_0{}).partition_A(cS);
1672 +
1673 + // have a mapping for each thread to coord
1674 + // find identical mapping to coords for the MMA
1675 + auto l = make_ordered_layout(make_shape(make_shape(_64{}, _2{}), make_shape(_16{}, TileShapeS{} / _32{})), make_stride(make_stride(_0{}, _3{}), make_stride(_1{}, _2{})));
1676 + auto sP_ = as_position_independent_swizzle_tensor(sP);
1677 + copy_aligned(tTR_rS_frag, sP_.compose(l)(threadIdx.x, _));
1678 +
1679 + // sum
1680 + row_sum *= correction_factor;
1681 +
1682 + static_assert(cute::is_same_v<ElementAcc, float>);
1683 + auto tTR_rAcc_float2 = recast<float2>(tTR_rAcc);
1684 + auto sums = make_tensor<float2>(_4{});
1685 + static_assert(size(tTR_rAcc_float2) % size(sums) == 0);
1686 + CUTLASS_PRAGMA_UNROLL
1687 + for (int i = 0; i < size(sums); i++) {
1688 + sums(i) = tTR_rAcc_float2(i);
1689 + }
1690 + CUTLASS_PRAGMA_UNROLL
1691 + for (int i = size(sums); i < size(tTR_rAcc_float2); i += size(sums)) {
1692 + CUTLASS_PRAGMA_UNROLL
1693 + for (int j = 0; j < size(sums); j++) {
1694 + cute::add(sums(j), sums(j), tTR_rAcc_float2(i + j));
1695 + }
1696 + }
1697 + CUTLASS_PRAGMA_UNROLL
1698 + for (int i = 1; i < size(sums); i *= 2) {
1699 + CUTLASS_PRAGMA_UNROLL
1700 + for (int j = 0; j < size(sums); j += 2*i) {
1701 + cute::add(sums(j), sums(j), sums(j+i));
1702 + }
1703 + }
1704 + row_sum += sums(0).x + sums(0).y;
1705 + }
1706 +
1707 +
1708 + CUTLASS_DEVICE void rescale(
1709 + ElementAcc correction_factor,
1710 + uint32_t tmem_o) {
1711 +
1712 + // for b2b gemm, do nothing
1713 +#ifndef B2B
1714 + auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{};
1715 + auto store_op = TMEM::tmem_load_to_store(load_op);
1716 +
1717 + TiledMmaPV tiled_mma_pv;
1718 +
1719 + Tensor tItI = partition_fragment_C(tiled_mma_pv, select<0,1>(TileShapePV{}));
1720 + tItI.data() = tmem_o;
1721 +
1722 + CUTE_STATIC_ASSERT_V(shape<1>(tItI) == _1{});
1723 + CUTE_STATIC_ASSERT_V(shape<2>(tItI) == _1{});
1724 + Tensor tAcc = tItI(make_coord(_,_),_0{},_0{});
1725 +
1726 + auto cta_tiler_pv = take<0,2>(typename CollectiveMmaPV::CtaShape_MNK{});
1727 + Tensor gO = make_tensor(make_gmem_ptr((ElementAcc*) nullptr), cta_tiler_pv, make_stride(0, 0));
1728 +
1729 + auto tiled_t2r = make_tmem_copy(load_op, tAcc);
1730 + auto tiled_r2t = make_tmem_copy(store_op, tAcc);
1731 + auto thread_idx = threadIdx.x % size(tiled_t2r);
1732 +
1733 + auto thread_t2r = tiled_t2r.get_slice(thread_idx);
1734 + auto thread_r2t = tiled_r2t.get_slice(thread_idx);
1735 + Tensor tTR_gO = thread_t2r.partition_D(gO);
1736 + Tensor tTR_rAcc = make_tensor<ElementAcc>(shape(tTR_gO));
1737 +
1738 + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc);
1739 +
1740 + // load o
1741 + copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
1742 +
1743 + // multiply by correction factor
1744 + float2 correction_factor_vec = make_float2(correction_factor, correction_factor);
1745 + CUTLASS_PRAGMA_UNROLL
1746 + for (int i = 0; i < size(tTR_rAcc); i += 2) {
1747 + float2 in = make_float2(tTR_rAcc(i + 0), tTR_rAcc(i + 1));
1748 + float2 out;
1749 + cute::mul(out, in, correction_factor_vec);
1750 + tTR_rAcc(i + 0) = out.x;
1751 + tTR_rAcc(i + 1) = out.y;
1752 + }
1753 +
1754 + // store o
1755 + copy(tiled_r2t, tTR_rAcc, tTR_tAcc);
1756 +#endif
1757 + }
1758 +
1759 +
1760 + template<class BlkCoord>
1761 + CUTLASS_DEVICE void epilogue(
1762 + ElementAcc& row_max,
1763 + ElementAcc& row_sum,
1764 + BlkCoord const& cta_coord,
1765 + ProblemShape const& problem_shape,
1766 + MainloopArguments const& mainloop_args,
1767 + EpilogueParams const& epilogue_args,
1768 + TensorStorage& shared_tensors,
1769 + uint32_t tmem_o,
1770 + int const& split_kv) {
1771 +
1772 + auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{};
1773 +
1774 + TiledMmaPV tiled_mma_pv;
1775 +
1776 + Tensor tItI = TiledMmaPV::make_fragment_C(partition_shape_C(TiledMmaPV{}, take<0, 2>(TileShapePV{})));
1777 + tItI.data() = tmem_o;
1778 +
1779 + CUTE_STATIC_ASSERT_V(shape<1>(tItI) == _1{});
1780 + CUTE_STATIC_ASSERT_V(shape<2>(tItI) == _1{});
1781 + Tensor tAcc = tItI(make_coord(_,_),_0{},_0{});
1782 +
1783 + auto [H, K, D, B] = problem_shape;
1784 + auto [D_latent, D_rope] = D;
1785 + if (epilogue_args.ptr_o_acc != nullptr) {
1786 + using ElementOutAcc = ElementAcc;
1787 + constexpr auto AlignmentOutAcc = 128 / cute::sizeof_bits_v<ElementOutAcc>;
1788 + Tensor mO = make_tensor(make_gmem_ptr(epilogue_args.ptr_o_acc + get<3>(cta_coord) * D_latent), make_shape(H, D_latent, B), epilogue_args.stride_o_acc);
1789 + auto cta_tiler_pv = take<0,2>(typename CollectiveMmaPV::CtaShape_MNK{});
1790 + Tensor gO = local_tile(mO, cta_tiler_pv, take<0,3>(cta_coord));
1791 +
1792 + auto tiled_t2r = make_tmem_copy(load_op, tAcc);
1793 + auto thread_idx = threadIdx.x % size(tiled_t2r);
1794 +
1795 + auto thread_t2r = tiled_t2r.get_slice(thread_idx);
1796 + Tensor tTR_gO = thread_t2r.partition_D(gO);
1797 + Tensor tTR_rAcc = make_tensor<ElementAcc>(shape(tTR_gO));
1798 +
1799 + Tensor tTR_rO_frag = make_tensor<ElementOutAcc>(shape(tTR_rAcc));
1800 + Tensor tTR_rO_src = recast<Array<ElementOutAcc, AlignmentOutAcc>>(coalesce(tTR_rO_frag));
1801 + Tensor tR2G_rO_dst = recast<Array<ElementOutAcc, AlignmentOutAcc>>(coalesce(tTR_gO));
1802 + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc);
1803 +
1804 + copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
1805 +
1806 + cutlass::epilogue::thread::LinearCombination<ElementOutAcc, 1, ElementAcc, ElementAcc, cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling> epilogue_op({epilogue_args.output_scale / row_sum});
1807 + CUTLASS_PRAGMA_UNROLL
1808 + for (int i = 0; i < size(tTR_rAcc); i++) {
1809 + tTR_rO_frag(i) = epilogue_op(tTR_rAcc(i));
1810 + }
1811 +
1812 + copy(tTR_rO_src, tR2G_rO_dst);
1813 +
1814 +#ifndef B2B
1815 +
1816 + // compute LSE
1817 + ElementAcc lse = cutlass::fast_log(row_sum) + mainloop_args.softmax_scale * row_max;
1818 +
1819 + // store LSE
1820 + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_args.ptr_lse_acc + H * get<3>(cta_coord)), make_shape(H, B), epilogue_args.stride_lse_acc);
1821 + Tensor gLSE = local_tile(mLSE, append<3>(cta_tiler_pv, _1{}), take<0,3>(cta_coord), Step<_1, Underscore, _1>{});
1822 + // for 2x2 dp, this must be conditional and the index is wrong
1823 + if (! kIs2Sm || (threadIdx.x < 64))
1824 + {
1825 + gLSE(threadIdx.x) = lse;
1826 + }
1827 + #endif
1828 + }
1829 + else {
1830 + Tensor mO = make_tensor(make_gmem_ptr(epilogue_args.ptr_o), make_shape(H, D_latent, B), epilogue_args.stride_o);
1831 + auto cta_tiler_pv = take<0,2>(typename CollectiveMmaPV::CtaShape_MNK{});
1832 + Tensor gO = local_tile(mO, cta_tiler_pv, take<0,3>(cta_coord));
1833 +
1834 + auto tiled_t2r = make_tmem_copy(load_op, tAcc);
1835 + auto thread_idx = threadIdx.x % size(tiled_t2r);
1836 +
1837 + auto thread_t2r = tiled_t2r.get_slice(thread_idx);
1838 + Tensor tTR_gO = thread_t2r.partition_D(gO);
1839 + Tensor tTR_rAcc = make_tensor<ElementAcc>(shape(tTR_gO));
1840 +
1841 + Tensor tTR_rO_frag = make_tensor<ElementOut>(shape(tTR_rAcc));
1842 + Tensor tTR_rO_src = recast<Array<ElementOut, AlignmentOut>>(coalesce(tTR_rO_frag));
1843 + Tensor tR2G_rO_dst = recast<Array<ElementOut, AlignmentOut>>(coalesce(tTR_gO));
1844 + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc);
1845 +
1846 + copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
1847 +
1848 + cutlass::epilogue::thread::LinearCombination<ElementOut, 1, ElementAcc, ElementAcc, cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling> epilogue_op({epilogue_args.output_scale / row_sum});
1849 + CUTLASS_PRAGMA_UNROLL
1850 + for (int i = 0; i < size(tTR_rAcc); i++) {
1851 + tTR_rO_frag(i) = epilogue_op(tTR_rAcc(i));
1852 + }
1853 +
1854 + copy(tTR_rO_src, tR2G_rO_dst);
1855 +
1856 +#ifndef B2B
1857 + if (epilogue_args.ptr_lse != nullptr) {
1858 + // compute LSE
1859 + ElementAcc lse = cutlass::fast_log(row_sum) + mainloop_args.softmax_scale * row_max;
1860 +
1861 + // store LSE
1862 + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_args.ptr_lse), make_shape(H, B), epilogue_args.stride_lse);
1863 + Tensor gLSE = local_tile(mLSE, append<3>(cta_tiler_pv, _1{}), take<0,3>(cta_coord), Step<_1, Underscore, _1>{});
1864 +
1865 + // for 2x2 dp, this must be conditional and the index is wrong
1866 + if (! kIs2Sm || (threadIdx.x < 64))
1867 + {
1868 + gLSE(threadIdx.x) = lse;
1869 + }
1870 + }
1871 +#endif
1872 + }
1873 + }
1874 +
1875 +
1876 + template<class CtaCoord>
1877 + CUTLASS_DEVICE void compute(
1878 + CtaCoord const& cta_coord,
1879 + ProblemShape const& problem_shape,
1880 + MainloopArguments const& mainloop_args,
1881 + EpilogueParams const& epilogue_args,
1882 + TensorStorage& shared_tensors,
1883 + PipelineS& pipeline_mma_s,
1884 + typename PipelineS::PipelineState& pipeline_mma_s_consumer_state,
1885 + PipelineP& pipeline_p_mma,
1886 + typename PipelineP::PipelineState& pipeline_p_mma_producer_state,
1887 + PipelineO& pipeline_mma_o,
1888 + typename PipelineO::PipelineState& pipeline_mma_o_consumer_state,
1889 + int const& split_kv) {
1890 +
1891 + auto [H, K, D, B] = problem_shape;
1892 +
1893 + int k_tile_total = ceil_div(K, TileShapeS{});
1894 + int k_tile_per_cta = ceil_div(k_tile_total, split_kv);
1895 + int k_index = get<3>(cta_coord) * k_tile_per_cta; // lower limit
1896 + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index);
1897 + if (k_tile_count == 0) {
1898 +
1899 + // if we return early, we have to make sure we release the load warp
1900 + cutlass::arch::NamedBarrier(
1901 + (kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp,
1902 + kNamedBarrierEpilogue
1903 + ).arrive();
1904 +
1905 + return;
1906 + }
1907 + int k_index_final = k_tile_total - 1;
1908 +
1909 + ElementAcc row_max = -std::numeric_limits<ElementAcc>::infinity();
1910 + ElementAcc row_sum = 0;
1911 + ElementAcc correction_factor = 1;
1912 +
1913 + pipeline_p_mma.producer_acquire(pipeline_p_mma_producer_state);
1914 + pipeline_mma_s.consumer_wait(pipeline_mma_s_consumer_state);
1915 +
1916 + auto dispatch_bool = [](bool b, auto fn) {
1917 + if (b) {
1918 + fn(cute::true_type{});
1919 + }
1920 + else {
1921 + fn(cute::false_type{});
1922 + }
1923 + };
1924 +
1925 + // softmax s0 -> p0
1926 + dispatch_bool(k_index == k_index_final, [&](auto is_last_tile) {
1927 + softmax(
1928 + is_last_tile,
1929 + row_max, row_sum, correction_factor,
1930 + problem_shape, mainloop_args, shared_tensors, k_index,
1931 + uint32_t(pipeline_mma_s_consumer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1),
1932 + pipeline_p_mma_producer_state.index()
1933 + );
1934 + });
1935 +
1936 + k_index += 1;
1937 +
1938 + cutlass::arch::fence_view_async_tmem_load();
1939 + cutlass::arch::fence_view_async_shared();
1940 + pipeline_mma_s.consumer_release(pipeline_mma_s_consumer_state);
1941 + ++pipeline_mma_s_consumer_state;
1942 + pipeline_p_mma.producer_commit(pipeline_p_mma_producer_state);
1943 + ++pipeline_p_mma_producer_state;
1944 +
1945 + k_tile_count -= 1;
1946 +
1947 + CUTLASS_PRAGMA_NO_UNROLL
1948 + while (k_tile_count > 0) {
1949 + pipeline_p_mma.producer_acquire(pipeline_p_mma_producer_state);
1950 + pipeline_mma_s.consumer_wait(pipeline_mma_s_consumer_state);
1951 +
1952 + // softmax s1 -> p1
1953 + dispatch_bool(k_index == k_index_final, [&](auto is_last_tile) {
1954 + softmax(
1955 + is_last_tile,
1956 + row_max, row_sum, correction_factor,
1957 + problem_shape, mainloop_args, shared_tensors, k_index,
1958 + uint32_t(pipeline_mma_s_consumer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1),
1959 + pipeline_p_mma_producer_state.index()
1960 + );
1961 + });
1962 +
1963 + cutlass::arch::fence_view_async_tmem_load();
1964 + cutlass::arch::fence_view_async_shared();
1965 + pipeline_mma_s.consumer_release(pipeline_mma_s_consumer_state);
1966 + ++pipeline_mma_s_consumer_state;
1967 + pipeline_p_mma.producer_commit(pipeline_p_mma_producer_state);
1968 + ++pipeline_p_mma_producer_state;
1969 +
1970 + pipeline_mma_o.consumer_wait(pipeline_mma_o_consumer_state);
1971 +
1972 + // rescale
1973 + CUTLASS_PRAGMA_UNROLL
1974 + for (int j = 0; j < IterationsPV_N; j++) {
1975 + rescale(correction_factor, uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO));
1976 + }
1977 +
1978 + cutlass::arch::fence_view_async_tmem_store();
1979 + pipeline_mma_o.consumer_release(pipeline_mma_o_consumer_state);
1980 + ++pipeline_mma_o_consumer_state;
1981 +
1982 + --k_tile_count;
1983 + k_index += 1;
1984 + }
1985 +
1986 + pipeline_mma_o.consumer_wait(pipeline_mma_o_consumer_state);
1987 +
1988 +#ifdef B2B
1989 + row_sum = 1;
1990 +#else
1991 + if constexpr (kWarpsInN > 1) {
1992 + // reduce row_sum if needed (for 2x2 dp)
1993 + shared_tensors.smem_exchange[threadIdx.x] = row_sum;
1994 + cutlass::arch::NamedBarrier(kNumComputeWarps*NumThreadsPerWarp, kNamedBarrierExchange).sync();
1995 + // (64, 2) shape
1996 + int peer_index = (threadIdx.x + 64) % 128;
1997 + row_sum += shared_tensors.smem_exchange[peer_index];
1998 + }
1999 +#endif
2000 +
2001 + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive();
2002 +
2003 + // epilogue
2004 + CUTLASS_PRAGMA_UNROLL
2005 + for (int j = 0; j < IterationsPV_N; j++) {
2006 + epilogue(
2007 + row_max, row_sum,
2008 + replace<1>(cta_coord, j), problem_shape,
2009 + mainloop_args, epilogue_args, shared_tensors,
2010 + uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO), split_kv
2011 + );
2012 + }
2013 +
2014 + cutlass::arch::fence_view_async_tmem_load();
2015 + pipeline_mma_o.consumer_release(pipeline_mma_o_consumer_state);
2016 + ++pipeline_mma_o_consumer_state;
2017 + }
2018 +
2019 +};
2020 +
2021 +///////////////////////////////////////////////////////////////////////////////
2022 +
2023 +} // namespace cutlass::fmha::kernel
csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp new
+165
@@ -0,0 +1,165 @@
1 +/***************************************************************************************************
2 + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
3 + *reserved. SPDX-License-Identifier: BSD-3-Clause
4 + *
5 + * Redistribution and use in source and binary forms, with or without
6 + * modification, are permitted provided that the following conditions are met:
7 + *
8 + * 1. Redistributions of source code must retain the above copyright notice,
9 + *this list of conditions and the following disclaimer.
10 + *
11 + * 2. Redistributions in binary form must reproduce the above copyright notice,
12 + * this list of conditions and the following disclaimer in the documentation
13 + * and/or other materials provided with the distribution.
14 + *
15 + * 3. Neither the name of the copyright holder nor the names of its
16 + * contributors may be used to endorse or promote products derived from
17 + * this software without specific prior written permission.
18 + *
19 + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20 + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21 + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22 + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
23 + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24 + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25 + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26 + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27 + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28 + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 + *POSSIBILITY OF SUCH DAMAGE.
30 + *
31 + **************************************************************************************************/
32 +/*
33 + * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
34 + * by Alcanderian JieXin Liang
35 + */
36 +
37 +// clang-format off
38 +#pragma once
39 +
40 +#include "cutlass/cutlass.h"
41 +#include "cutlass/fast_math.h"
42 +#include "cutlass/kernel_hardware_info.h"
43 +
44 +namespace cutlass::fmha::kernel {
45 +
46 +////////////////////////////////////////////////////////////////////////////////
47 +
48 +struct Sm100MlaIndividualTileScheduler {
49 +
50 + struct Params {
51 + dim3 grid;
52 + };
53 +
54 + bool valid_ = true;
55 +
56 + CUTLASS_DEVICE
57 + Sm100MlaIndividualTileScheduler(Params const&) {}
58 +
59 + template<class ProblemShape, class ClusterShape>
60 + static Params to_underlying_arguments(
61 + ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
62 + ClusterShape const& cluster_shape, int const& split_kv) {
63 + using namespace cute;
64 + dim3 grid(get<0>(cluster_shape), get<3>(problem_shape) /* Batch */, split_kv /*Maximum Split KV*/);
65 + return Params{ grid };
66 + }
67 +
68 + static dim3 get_grid_shape(Params const& params) {
69 + return params.grid;
70 + }
71 +
72 + CUTLASS_DEVICE
73 + bool is_valid() {
74 + return valid_;
75 + }
76 +
77 + CUTLASS_DEVICE
78 + auto get_block_coord() {
79 + using namespace cute;
80 + return make_coord(blockIdx.x, _0{}, blockIdx.y, blockIdx.z);
81 + }
82 +
83 + CUTLASS_DEVICE
84 + Sm100MlaIndividualTileScheduler& operator++() {
85 + valid_ = false;
86 + return *this;
87 + }
88 +};
89 +
90 +////////////////////////////////////////////////////////////////////////////////
91 +
92 +struct Sm100MlaPersistentTileScheduler {
93 +
94 + struct Params {
95 + int num_blocks;
96 + FastDivmod divmod_m_block;
97 + FastDivmod divmod_b;
98 + FastDivmod divmod_split_kv;
99 + KernelHardwareInfo hw_info;
100 + };
101 +
102 + int block_idx = 0;
103 + Params params;
104 +
105 + CUTLASS_DEVICE
106 + Sm100MlaPersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {}
107 +
108 + template<class ProblemShape, class ClusterShape>
109 + static Params to_underlying_arguments(
110 + ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
111 + ClusterShape const& cluster_shape, int const& split_kv) {
112 + using namespace cute;
113 + // Get SM count if needed, otherwise use user supplied SM count
114 + int sm_count = hw_info.sm_count;
115 + if (sm_count <= 1 || sm_count % size<0>(cluster_shape) != 0) {
116 + CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
117 + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
118 + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
119 + }
120 +
121 + CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
122 + hw_info.sm_count = sm_count;
123 +
124 + int num_m_blocks = size<0>(cluster_shape);
125 + int num_blocks = num_m_blocks * get<3>(problem_shape) /* Batch */;
126 + num_blocks *= split_kv; /* Maximum Split KV*/
127 +
128 + return Params {
129 + num_blocks,
130 + { num_m_blocks}, { get<3>(problem_shape) }, {split_kv},
131 + hw_info
132 + };
133 + }
134 +
135 + static dim3 get_grid_shape(Params const& params) {
136 + dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1);
137 + return grid;
138 + }
139 +
140 + CUTLASS_DEVICE
141 + bool is_valid() {
142 + return block_idx < params.num_blocks;
143 + }
144 +
145 + CUTLASS_DEVICE
146 + auto get_block_coord() {
147 + using namespace cute;
148 + int block_decode = block_idx;
149 + int m_block, bidb, n_split_kv;
150 + params.divmod_m_block(block_decode, m_block, block_decode);
151 + params.divmod_b(block_decode, bidb, block_decode);
152 + params.divmod_split_kv(block_decode, n_split_kv, block_decode);
153 + return make_coord(m_block, _0{}, bidb, n_split_kv);
154 + }
155 +
156 + CUTLASS_DEVICE
157 + Sm100MlaPersistentTileScheduler& operator++() {
158 + block_idx += gridDim.x;
159 + return *this;
160 + }
161 +};
162 +
163 +////////////////////////////////////////////////////////////////////////////////
164 +
165 +} // namespace cutlass::fmha::kernel
csrc/attention/mla/sm100_cutlass_mla_kernel.cu new
+273
@@ -0,0 +1,273 @@
1 +/*
2 +Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3 +Copyright 2025 SGLang Team. All Rights Reserved.
4 +
5 +Licensed under the Apache License, Version 2.0 (the "License");
6 +you may not use this file except in compliance with the License.
7 +You may obtain a copy of the License at
8 +
9 + http://www.apache.org/licenses/LICENSE-2.0
10 +
11 +Unless required by applicable law or agreed to in writing, software
12 +distributed under the License is distributed on an "AS IS" BASIS,
13 +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 +See the License for the specific language governing permissions and
15 +limitations under the License.
16 +==============================================================================*/
17 +/*
18 + * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
19 + * by Alcanderian JieXin Liang
20 + */
21 +
22 +#include <ATen/cuda/CUDAContext.h>
23 +#include <c10/cuda/CUDAGuard.h>
24 +#include <cutlass/cutlass.h>
25 +#include <cutlass/kernel_hardware_info.h>
26 +#include <torch/all.h>
27 +
28 +#include <cute/tensor.hpp>
29 +#include <iostream>
30 +
31 +#include "cutlass_sm100_mla/device/sm100_mla.hpp"
32 +#include "cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp"
33 +
34 +// clang-format off
35 +#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
36 +void sm100_cutlass_mla_decode(
37 + torch::Tensor const& out,
38 + torch::Tensor const& q_nope,
39 + torch::Tensor const& q_pe,
40 + torch::Tensor const& kv_c_and_k_pe_cache,
41 + torch::Tensor const& seq_lens,
42 + torch::Tensor const& page_table,
43 + torch::Tensor const& workspace,
44 + int64_t num_kv_splits) {
45 + TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
46 +}
47 +int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
48 + TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
49 +}
50 +#else
51 +
52 +#define CUTLASS_CHECK(status) \
53 + { \
54 + cutlass::Status error = status; \
55 + TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
56 + }
57 +
58 +using namespace cute;
59 +using namespace cutlass::fmha::kernel;
60 +
61 +template <bool v>
62 +struct IsPersistent {
63 + static const bool value = v;
64 +};
65 +
66 +template <typename T, bool IsPaged128, typename PersistenceOption = IsPersistent<true>>
67 +struct MlaSm100 {
68 + using Element = T;
69 + using ElementAcc = float;
70 + using ElementOut = T;
71 +
72 + using TileShape = Shape<_128, _128, Shape<_512, _64>>;
73 + using TileShapeH = cute::tuple_element_t<0, TileShape>;
74 + using TileShapeD = cute::tuple_element_t<2, TileShape>;
75 +
76 + // H K (D_latent D_rope) B
77 + using ProblemShape = cute::tuple<TileShapeH, int, TileShapeD, int>;
78 +
79 + using StrideQ = cute::tuple<int64_t, _1, int64_t>; // H D B
80 + using StrideK = cute::tuple<int64_t, _1, int64_t>; // K D B
81 + using StrideO = StrideK; // H D B
82 + using StrideLSE = cute::tuple<_1, int>; // H B
83 +
84 + using TileScheduler =
85 + std::conditional_t<PersistenceOption::value, Sm100MlaPersistentTileScheduler, Sm100MlaIndividualTileScheduler>;
86 +
87 + using FmhaKernel = cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized<
88 + TileShape,
89 + Element,
90 + ElementAcc,
91 + ElementOut,
92 + ElementAcc,
93 + TileScheduler,
94 + /*kIsCpAsync=*/!IsPaged128>;
95 + using Fmha = cutlass::fmha::device::MLA<FmhaKernel>;
96 +};
97 +
98 +template <typename T>
99 +typename T::Fmha::Arguments args_from_options(
100 + at::Tensor const& out,
101 + at::Tensor const& q_nope,
102 + at::Tensor const& q_pe,
103 + at::Tensor const& kv_c_and_k_pe_cache,
104 + at::Tensor const& seq_lens,
105 + at::Tensor const& page_table,
106 + double sm_scale,
107 + int64_t num_kv_splits) {
108 + cutlass::KernelHardwareInfo hw_info;
109 + hw_info.device_id = q_nope.device().index();
110 + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
111 +
112 + int batches = q_nope.sizes()[0];
113 + int page_count_per_seq = page_table.sizes()[1];
114 + int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
115 + int page_size = kv_c_and_k_pe_cache.sizes()[1];
116 + int max_seq_len = page_size * page_count_per_seq;
117 + using TileShapeH = typename T::TileShapeH;
118 + using TileShapeD = typename T::TileShapeD;
119 + auto problem_shape = cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches);
120 +
121 + auto [H, K, D, B] = problem_shape;
122 + auto [D_latent, D_rope] = D;
123 +
124 + float scale = float(sm_scale);
125 +
126 + using StrideQ = typename T::StrideQ;
127 + using StrideK = typename T::StrideK;
128 + using StrideO = typename T::StrideO;
129 + using StrideLSE = typename T::StrideLSE;
130 +
131 + StrideQ stride_Q_nope = cute::make_tuple(
132 + static_cast<int64_t>(q_nope.stride(1)), _1{}, static_cast<int64_t>(q_nope.stride(0)));
133 + StrideQ stride_Q_pe = cute::make_tuple(
134 + static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
135 +
136 + StrideK stride_C = cute::make_tuple(
137 + static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
138 + StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
139 + StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
140 + StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
141 +
142 + using Element = typename T::Element;
143 + using ElementOut = typename T::ElementOut;
144 + using ElementAcc = typename T::ElementAcc;
145 + auto Q_nope_ptr = static_cast<Element*>(q_nope.data_ptr());
146 + auto Q_pe_ptr = static_cast<Element*>(q_pe.data_ptr());
147 + auto C_ptr = static_cast<Element*>(kv_c_and_k_pe_cache.data_ptr());
148 + typename T::Fmha::Arguments arguments{
149 + problem_shape,
150 + {scale,
151 + Q_nope_ptr,
152 + stride_Q_nope,
153 + Q_pe_ptr,
154 + stride_Q_pe,
155 + C_ptr,
156 + stride_C,
157 + C_ptr + D_latent,
158 + stride_C,
159 + static_cast<int*>(seq_lens.data_ptr()),
160 + static_cast<int*>(page_table.data_ptr()),
161 + stride_PT,
162 + page_count_total,
163 + page_size},
164 + {static_cast<ElementOut*>(out.data_ptr()), stride_O, static_cast<ElementAcc*>(nullptr), stride_LSE},
165 + hw_info,
166 + // TODO(trevor-m): Change split_kv back to -1 when
167 + // https://github.com/NVIDIA/cutlass/issues/2274 is fixed. Split_kv=1 will
168 + // perform worse with larger context length and smaller batch sizes.
169 + num_kv_splits, // split_kv
170 + nullptr, // is_var_split_kv
171 + };
172 + // TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute
173 + // split_kv automatically based on batch size and sequence length to balance
174 + // workload across available SMs. Consider using var_split_kv for manual
175 + // control if needed.
176 + T::Fmha::set_split_kv(arguments);
177 + return arguments;
178 +}
179 +
180 +template <typename Element, bool IsPaged128, typename PersistenceOption>
181 +void runMla(
182 + at::Tensor const& out,
183 + at::Tensor const& q_nope,
184 + at::Tensor const& q_pe,
185 + at::Tensor const& kv_c_and_k_pe_cache,
186 + at::Tensor const& seq_lens,
187 + at::Tensor const& page_table,
188 + at::Tensor const& workspace,
189 + double sm_scale,
190 + int64_t num_kv_splits,
191 + cudaStream_t stream) {
192 + using MlaSm100Type = MlaSm100<Element, IsPaged128, PersistenceOption>;
193 + typename MlaSm100Type::Fmha fmha;
194 + auto arguments = args_from_options<MlaSm100Type>(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, sm_scale, num_kv_splits);
195 +
196 + CUTLASS_CHECK(fmha.can_implement(arguments));
197 +
198 + CUTLASS_CHECK(fmha.initialize(arguments, workspace.data_ptr(), stream));
199 +
200 + CUTLASS_CHECK(fmha.run(arguments, workspace.data_ptr(), stream));
201 +}
202 +
203 +#define DISPATCH_BOOL(expr, const_expr, ...) \
204 + [&]() -> bool { \
205 + if (expr) { \
206 + constexpr bool const_expr = true; \
207 + return __VA_ARGS__(); \
208 + } else { \
209 + constexpr bool const_expr = false; \
210 + return __VA_ARGS__(); \
211 + } \
212 + }()
213 +
214 +void sm100_cutlass_mla_decode(
215 + torch::Tensor const& out,
216 + torch::Tensor const& q_nope,
217 + torch::Tensor const& q_pe,
218 + torch::Tensor const& kv_c_and_k_pe_cache,
219 + torch::Tensor const& seq_lens,
220 + torch::Tensor const& page_table,
221 + torch::Tensor const& workspace,
222 + double sm_scale,
223 + int64_t num_kv_splits) {
224 + auto in_dtype = q_nope.dtype();
225 + at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
226 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
227 + const int page_size = kv_c_and_k_pe_cache.sizes()[1];
228 +
229 + // NOTE(alcanderian): IsPersistent has bug with manual split_kv.
230 + // Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
231 + // Maybe per batch split kv will fix this.
232 + DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
233 + DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
234 + if (in_dtype == at::ScalarType::Half) {
235 + runMla<cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
236 + out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
237 + } else if (in_dtype == at::ScalarType::BFloat16) {
238 + runMla<cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
239 + out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
240 + } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
241 + runMla<cutlass::float_e4m3_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
242 + out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
243 + } else {
244 + TORCH_CHECK(false, "Unsupported input data type of MLA");
245 + }
246 + return true;
247 + });
248 + return true;
249 + });
250 +}
251 +
252 +int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
253 + // Workspace size depends on ElementAcc and ElementLSE (same as ElementAcc)
254 + // which are float, so Element type here doesn't matter.
255 + using MlaSm100Type = MlaSm100<cutlass::half_t, true>;
256 +
257 + // Get split kv. Requires problem shape and sm_count only.
258 + typename MlaSm100Type::Fmha::Arguments arguments;
259 + using TileShapeH = typename MlaSm100Type::TileShapeH;
260 + using TileShapeD = typename MlaSm100Type::TileShapeD;
261 + arguments.problem_shape =
262 + cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
263 + // Assumes device 0 when getting sm_count.
264 + arguments.hw_info.sm_count =
265 + sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
266 + arguments.split_kv = num_kv_splits;
267 + MlaSm100Type::Fmha::set_split_kv(arguments);
268 +
269 + return MlaSm100Type::Fmha::get_workspace_size(arguments);
270 +}
271 +
272 +#endif
273 +// clang-format on
csrc/ops.h
+13
@@ -167,6 +167,19 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
167 torch::Tensor const& seq_lens,
168 torch::Tensor const& page_table, double scale);
169
170 +void sm100_cutlass_mla_decode(
171 + torch::Tensor const& out, torch::Tensor const& q_nope,
172 + torch::Tensor const& q_pe, torch::Tensor const& kv_c_and_k_pe_cache,
173 + torch::Tensor const& seq_lens, torch::Tensor const& page_table,
174 + torch::Tensor const& workspace, double sm_scale,
175 + int64_t num_kv_splits =
176 + 1 /* Set to 1 to avoid cuda_graph issue by default. */);
177 +
178 +int64_t sm100_cutlass_mla_get_workspace_size(
179 + int64_t max_seq_len, int64_t num_batches, int64_t sm_count = 0,
180 + int64_t num_kv_splits =
181 + 1 /* Set to 1 to avoid cuda_graph issue by default. */);
182 +
183 torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
184
185 #ifndef USE_ROCM
csrc/torch_bindings.cpp
+17
@@ -514,6 +514,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
514 " Tensor page_table, float scale) -> ()");
515 ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
516
517 + // SM100 CUTLASS MLA decode
518 + ops.def(
519 + "sm100_cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe,"
520 + " Tensor kv_c_and_k_pe_cache, Tensor seq_lens,"
521 + " Tensor page_table, Tensor workspace, float "
522 + "scale,"
523 + " int num_kv_splits) -> ()");
524 + ops.impl("sm100_cutlass_mla_decode", torch::kCUDA, &sm100_cutlass_mla_decode);
525 +
526 + // SM100 CUTLASS MLA workspace
527 + ops.def(
528 + "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
529 + " int sm_count, int num_kv_splits) "
530 + "-> int");
531 + ops.impl("sm100_cutlass_mla_get_workspace_size",
532 + &sm100_cutlass_mla_get_workspace_size);
533 +
534 // Compute NVFP4 block quantized tensor.
535 ops.def(
536 "scaled_fp4_quant(Tensor! output, Tensor input,"
vllm/_custom_ops.py
+20
@@ -1843,6 +1843,26 @@ def cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,
1843 return out
1844
1845
1846 +def sm100_cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,
1847 + q_pe: torch.Tensor,
1848 + kv_c_and_k_pe_cache: torch.Tensor,
1849 + seq_lens: torch.Tensor, page_table: torch.Tensor,
1850 + workspace: torch.Tensor, scale: float,
1851 + num_kv_splits: int) -> torch.Tensor:
1852 + torch.ops._C.sm100_cutlass_mla_decode(out, q_nope, q_pe,
1853 + kv_c_and_k_pe_cache, seq_lens,
1854 + page_table, workspace, scale,
1855 + num_kv_splits)
1856 + return out
1857 +
1858 +
1859 +def sm100_cutlass_mla_get_workspace_size(max_seq_len: int, num_batches: int,
1860 + sm_count: int,
1861 + num_kv_splits: int) -> int:
1862 + return torch.ops._C.sm100_cutlass_mla_get_workspace_size(
1863 + max_seq_len, num_batches, sm_count, num_kv_splits)
1864 +
1865 +
1866 if hasattr(torch.ops._C, "weight_packed_linear"):
1867
1868 @register_fake("_C::weight_packed_linear")
vllm/platforms/cuda.py
+7
@@ -166,6 +166,13 @@ class CudaPlatformBase(Platform):
166 logger.info(
167 "Forcing kv cache block size to 64 for FlashMLA backend.")
168
169 + use_cutlass_mla = (envs.VLLM_ATTENTION_BACKEND is not None \
170 + and envs.VLLM_ATTENTION_BACKEND == "CUTLASS_MLA_VLLM_V1")
171 + if use_cutlass_mla and cache_config.block_size != 128:
172 + cache_config.block_size = 128
173 + logger.info("Forcing kv cache block size to 128 for "
174 + "CUTLASS_MLA_VLLM_V1 backend.")
175 +
176 compilation_config = vllm_config.compilation_config
177 if (envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput"
178 and parallel_config.data_parallel_size > 1
vllm/v1/attention/backends/mla/common.py
+5
@@ -333,6 +333,9 @@ class MLACommonMetadata(Generic[D]):
333 # |-------------------- seq_len ---------------------|
334 # |-- query_len ---|
335
336 + num_reqs: int
337 + max_query_len: int
338 +
339 num_actual_tokens: int # Number of tokens excluding padding.
340 query_start_loc: torch.Tensor
341 slot_mapping: torch.Tensor
@@ -716,6 +719,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
719 )
720
721 attn_metadata = self.metadata_cls(
722 + num_reqs=common_attn_metadata.num_reqs,
723 + max_query_len=common_attn_metadata.max_query_len,
724 num_actual_tokens=num_actual_tokens,
725 query_start_loc=query_start_loc,
726 slot_mapping=slot_mapping,
vllm/v1/attention/backends/mla/cutlass_mla.py
+183 -1
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import os
5 from typing import Any, Optional
6
7 import torch
@@ -27,6 +28,41 @@ class CutlassMLABackend(MLACommonBackend):
28 return CutlassMLAImpl
29
30
31 +class SM100Workspace:
32 +
33 + def __init__(self, initial_workspace_size):
34 + self._workspace_buf = torch.empty(initial_workspace_size,
35 + device="cuda",
36 + dtype=torch.uint8)
37 +
38 + self._block_size = 128 # Forced to 128
39 +
40 + # Pre-compute sm_count to avoid recomputing it. Use device 0 as a proxy
41 + # (assumes all devices are similar)
42 + properties = torch.cuda.get_device_properties(torch.device("cuda:0"))
43 + self._sm_count = properties.multi_processor_count
44 +
45 + def get_buf(self):
46 + return self._workspace_buf
47 +
48 + def ensure_size(self, attn_metadata: MLACommonMetadata,
49 + num_kv_splits: int):
50 + batch_size = attn_metadata.num_reqs
51 + max_seq_len = attn_metadata.max_query_len
52 +
53 + workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
54 + max_seq_len * self._block_size,
55 + batch_size,
56 + self._sm_count,
57 + num_kv_splits=num_kv_splits)
58 +
59 + if self._workspace_buf.shape[0] < workspace_size:
60 + self._workspace_buf.resize_(workspace_size)
61 +
62 +
63 +g_sm100_workspace = SM100Workspace(128 * 1024 * 1024) # 128MB
64 +
65 +
66 class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
67
68 def __init__(
@@ -68,7 +104,137 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
104 raise NotImplementedError(
105 "CutlassMLA V1 with FP8 KV cache not yet supported")
106
71 - def _forward_decode(
107 + self._use_old_cutlass_mla = False
108 + force_old_cutlass = os.environ.get("FORCE_OLD_CUTLASS_MLA", None)
109 + if force_old_cutlass:
110 + logger.warning("Forcing old cutlass mla kernel")
111 + self._use_old_cutlass_mla = True
112 +
113 + # TODO: Currently, num_kv_splits is limited to 16 to avoid hanging
114 + # issues. In case the code hangs, use:
115 + # FORCE_NUM_KV_SPLITS=1
116 + force_num_kv_splits = os.environ.get("FORCE_NUM_KV_SPLITS", None)
117 + if force_num_kv_splits:
118 + logger.warning("Forcing num_kv_splits to %d",
119 + int(force_num_kv_splits))
120 + self._num_kv_splits = int(force_num_kv_splits)
121 + else:
122 + self._num_kv_splits = -1 # => Auto-detect
123 +
124 + # Share workspace buffer across all executions
125 + self._workspace = g_sm100_workspace
126 +
127 + def _sm100_cutlass_mla_decode(
128 + self,
129 + q_nope: torch.Tensor,
130 + q_pe: torch.Tensor,
131 + kv_c_and_k_pe_cache: torch.Tensor,
132 + seq_lens: torch.Tensor,
133 + page_table: torch.Tensor,
134 + workspace: torch.Tensor,
135 + sm_scale: float,
136 + num_kv_splits: int,
137 + ) -> torch.Tensor:
138 + assert (q_nope.ndim == 3
139 + ), f"q_nope must be a 3D tensor, but got {q_nope.ndim}"
140 + assert (
141 + q_pe.ndim == 3), f"q_pe must be a 3D tensor, but got {q_pe.ndim}"
142 + assert (
143 + kv_c_and_k_pe_cache.ndim == 3
144 + ), "kv_c_and_k_pe_cache must be a 3D tensor, but got {}".format(
145 + kv_c_and_k_pe_cache.ndim)
146 +
147 + B_q, H, D_q_nope = q_nope.shape
148 + B_q_2, H_2, D_q_pe = q_pe.shape
149 + assert (B_q == B_q_2) and (H == H_2)
150 +
151 + _, PAGE_SIZE, D_ckv = kv_c_and_k_pe_cache.shape
152 +
153 + D_latent = 512
154 + D_rope = 64
155 + assert D_q_nope == D_latent
156 + assert D_q_pe == D_rope
157 + assert D_ckv == D_latent + D_rope
158 +
159 + MAX_HEADS = 128
160 + assert H <= MAX_HEADS, f"H must be <= {MAX_HEADS}, but got {H}"
161 + if H < MAX_HEADS:
162 + q_nope_padded = q_nope.new_empty((B_q, MAX_HEADS, D_q_nope))
163 + q_nope_padded[:, :H] = q_nope
164 + q_nope = q_nope_padded
165 +
166 + q_pe_padded = q_pe.new_empty((B_q, MAX_HEADS, D_q_pe))
167 + q_pe_padded[:, :H] = q_pe
168 + q_pe = q_pe_padded
169 +
170 + assert len(page_table.shape) == 2
171 + B_block_table, block_num = page_table.shape
172 + assert B_block_table == B_q
173 + assert (block_num
174 + > 0), f"block num must be greater than 0, got {block_num}"
175 + assert block_num % (128 / PAGE_SIZE) == 0
176 +
177 + # TODO(kaixih@nvidia): support fp8
178 + assert q_nope.dtype in (
179 + torch.float16,
180 + torch.bfloat16,
181 + ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}."
182 + assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype
183 + assert (
184 + seq_lens.dtype == torch.int32
185 + ), f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}."
186 + assert (
187 + page_table.dtype == torch.int32
188 + ), f"page_table.dtype needs to be int32 but got {page_table.dtype}."
189 +
190 + out = q_nope.new_empty((B_q, MAX_HEADS, D_latent))
191 +
192 + ops.sm100_cutlass_mla_decode(
193 + out,
194 + q_nope,
195 + q_pe,
196 + kv_c_and_k_pe_cache,
197 + seq_lens,
198 + page_table,
199 + workspace,
200 + sm_scale,
201 + num_kv_splits,
202 + )
203 + return out[:, :H].contiguous()
204 +
205 + def _sm100_forward_decode(
206 + self,
207 + q_nope: torch.Tensor,
208 + q_pe: torch.Tensor,
209 + kv_c_and_k_pe_cache: torch.Tensor,
210 + attn_metadata: MLACommonMetadata,
211 + ) -> torch.Tensor:
212 + assert kv_c_and_k_pe_cache.numel() > 0
213 + assert attn_metadata.decode is not None
214 +
215 + if self.kv_cache_dtype.startswith("fp8"):
216 + raise NotImplementedError("FP8 Cutlass MLA not yet supported")
217 +
218 + # Adjust workspace size (if necessary)
219 + self._workspace.ensure_size(attn_metadata, self._num_kv_splits)
220 +
221 + # Run MLA
222 + # Clone q_nope and q_pe to make sure strides computation is correct.
223 + # TODO: Check if we really need it
224 + q_nope = q_nope.clone()
225 + q_pe = q_pe.clone()
226 +
227 + o = self._sm100_cutlass_mla_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
228 + attn_metadata.decode.seq_lens,
229 + attn_metadata.decode.block_table,
230 + self._workspace.get_buf(),
231 + self.scale, self._num_kv_splits)
232 +
233 + return self._v_up_proj(o)
234 +
235 + # TODO: Currently we leave it here only for backup in case something is
236 + # wrong with the new SM100 CUTLASS MLA kernel
237 + def _old_forward_decode(
238 self,
239 q_nope: torch.Tensor,
240 q_pe: torch.Tensor,
@@ -97,3 +263,19 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
263 attn_metadata.decode.block_table, self.scale)
264
265 return self._v_up_proj(o)
266 +
267 + def _forward_decode(
268 + self,
269 + q_nope: torch.Tensor,
270 + q_pe: torch.Tensor,
271 + kv_c_and_k_pe_cache: torch.Tensor,
272 + attn_metadata: MLACommonMetadata,
273 + ) -> torch.Tensor:
274 + if self._use_old_cutlass_mla:
275 + # TODO: Remove the old cutlass MLA kernel after more extensive
276 + # testing
277 + return self._old_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
278 + attn_metadata)
279 +
280 + return self._sm100_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
281 + attn_metadata)