Skip to content
This repository was archived by the owner on May 13, 2026. It is now read-only.

Commit b492a3a

Browse files
guocuimiguocuimi
committed
move local related parameters out of block
1 parent 4327f74 commit b492a3a

7 files changed

Lines changed: 95 additions & 79 deletions

File tree

‎src/kernels/attention/collective/sm120_collective_fmha_mainloop_ws.cuh‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -179,16 +179,35 @@ struct Sm120CollectiveFMhaWs {
179179

180180
// load Q/K/V from gmem to smem
181181
template <class Block>
182-
CUTE_DEVICE void load(const Block& block,
182+
CUTE_DEVICE void load(const Params& params,
183+
const Block& block,
183184
int tidx,
184185
PipelineQ& q_pipeline,
185186
typename PipelineQ::PipelineState& q_state,
186187
PipelineKV& kv_pipeline,
187188
typename PipelineKV::PipelineState& kv_state,
188189
TensorStorage& ss) {
190+
if (!block.is_valid()) {
191+
// skip invalid block
192+
return;
193+
}
194+
const auto [n_block_min, n_block_max] =
195+
block.template get_kv_blocks<LOCAL>(params.sliding_window);
196+
if (n_block_min >= n_block_max) {
197+
return; // no kv blocks to process
198+
}
199+
189200
// forward to the load implementation
190201
Load load;
191-
load(block, tidx, q_pipeline, q_state, kv_pipeline, kv_state, ss);
202+
load(block,
203+
tidx,
204+
n_block_min,
205+
n_block_max,
206+
q_pipeline,
207+
q_state,
208+
kv_pipeline,
209+
kv_state,
210+
ss);
192211
}
193212

194213
template <class Block, class FrgTensor, class PipelineQ, class PipelineKV>
@@ -212,7 +231,8 @@ struct Sm120CollectiveFMhaWs {
212231
return;
213232
}
214233

215-
const auto [n_block_min, n_block_max] = block.get_kv_blocks();
234+
const auto [n_block_min, n_block_max] =
235+
block.template get_kv_blocks<LOCAL>(params.sliding_window);
216236
if (n_block_min >= n_block_max) {
217237
return; // no kv blocks to process
218238
}

‎src/kernels/attention/collective/sm120_collective_load_cpasync_ws.cuh‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,23 +39,15 @@ struct Sm120CollectiveLoadCpAsyncWs {
3939
template <class Block>
4040
CUTE_DEVICE void operator()(const Block& block,
4141
int tidx,
42+
int n_block_min,
43+
int n_block_max,
4244
PipelineQ& q_pipeline,
4345
typename PipelineQ::PipelineState& q_state,
4446
PipelineKV& kv_pipeline,
4547
typename PipelineKV::PipelineState& kv_state,
4648
TensorStorage& ss) {
4749
static constexpr int kStages = size<2>(SmemLayoutK{});
4850

49-
if (!block.is_valid()) {
50-
// skip invalid block
51-
return;
52-
}
53-
54-
const auto [n_block_min, n_block_max] = block.get_kv_blocks();
55-
if (n_block_min >= n_block_max) {
56-
return; // no kv blocks to process
57-
}
58-
5951
// (M, N, K)
6052
const auto residue_mnk = block.get_residue_mnk();
6153

‎src/kernels/attention/common/fmha_block.h‎

Lines changed: 48 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,14 @@ namespace llm {
1212
using namespace cute;
1313

1414
// AttentionTile specialization for AttentionParams
15-
template <typename TileShape, // (BLK_M, BLK_N, BLK_K)
16-
typename Element, // Element type
17-
typename StrideQ, // (Q, D, ((KH, G), B))
18-
typename StrideK, // (K, D, ((KH, _0), B))
19-
typename StrideV, // (V, D, ((KH, _0), B))
20-
typename StrideO, // (Q, Q, ((KH, G), B))
21-
bool kLocal>
15+
template <typename TileShape, // (BLK_M, BLK_N, BLK_K)
16+
typename BlocKCoord, // (m_block_idx, ((kv_head_idx, _0), batch_idx))
17+
typename Element, // Element type
18+
typename StrideQ, // (Q, D, ((KH, G), B))
19+
typename StrideK, // (K, D, ((KH, _0), B))
20+
typename StrideV, // (V, D, ((KH, _0), B))
21+
typename StrideO // (Q, Q, ((KH, G), B))
22+
>
2223
struct FmhaBlock {
2324
// Host side parameters
2425
struct Arguments {
@@ -31,8 +32,6 @@ struct FmhaBlock {
3132
StrideK k_stride;
3233
StrideV v_stride;
3334
StrideO o_stride;
34-
35-
int sliding_window = -1; // -1 means no sliding window
3635
};
3736

3837
// Device side parameters
@@ -47,8 +46,6 @@ struct FmhaBlock {
4746
StrideV v_stride;
4847
StrideO o_stride;
4948

50-
int sliding_window;
51-
5249
// Parameters from problem shape
5350
int batch_size;
5451
int q_len;
@@ -59,18 +56,6 @@ struct FmhaBlock {
5956
FastDivmod group_size;
6057
};
6158

62-
// using StrideK = ...;
63-
64-
// using TMA_K = decltype(make_tma_copy(
65-
// GmemTiledCopy{}, // TMA_COPY
66-
// make_tensor(static_cast<InternalElementA const*>(nullptr),
67-
// repeat_like(StrideK{}, int32_t(0)), StrideK{}),
68-
// SmemLayoutK{}(_,_,_0{})));
69-
70-
// Tensor tensor_k = make_tensor(ptr_k, make_layout(make_shape(M,K,L),
71-
// args.stride_k)); auto tma_load_k = make_tma_copy(SM90_TMA_LOAD{},
72-
// gtensor_k, SmemLayoutK{}(_,_,_0{}));
73-
7459
template <class ProblemShape>
7560
static Params to_underlying_arguments(const ProblemShape& problem_shape,
7661
const Arguments& args,
@@ -93,7 +78,6 @@ struct FmhaBlock {
9378
.k_stride = args.k_stride,
9479
.v_stride = args.v_stride,
9580
.o_stride = args.o_stride,
96-
.sliding_window = args.sliding_window,
9781
.batch_size = batch_size,
9882
.q_len = q_len,
9983
.kv_len = kv_len,
@@ -113,9 +97,6 @@ struct FmhaBlock {
11397

11498
// hold a reference to the parameters and block coordination
11599
const Params& params_;
116-
// TODO: pass in as parameter is better
117-
// (m_block_idx, ((kv_head_idx, _0), batch_idx))
118-
using BlocKCoord = Coord<int, Coord<Coord<int, _0>, int>>;
119100
const BlocKCoord& blk_coord_;
120101

121102
// derived parameters to avoid recomputation
@@ -161,9 +142,9 @@ struct FmhaBlock {
161142
// returns (m_block_idx, ((kv_head_idx, _0), batch_idx))
162143
CUTE_HOST_DEVICE const auto& get_block_coord() const { return blk_coord_; }
163144

164-
// TODO: pass in kLocal and sliding_window instead
165145
// returns kv block range: (n_block_min, n_block_max]
166-
CUTE_HOST_DEVICE auto get_kv_blocks() const {
146+
template <bool kLocal>
147+
CUTE_HOST_DEVICE auto get_kv_blocks(int sliding_window) const {
167148
static constexpr int kBlockM = get<0>(TileShape{});
168149
static constexpr int kBlockN = get<1>(TileShape{});
169150

@@ -176,7 +157,7 @@ struct FmhaBlock {
176157
const int n_block_max = cute::ceil_div(kv_idx_max, kBlockN);
177158

178159
if constexpr (kLocal) {
179-
const int kv_idx_min = std::max(0, diagonal - params_.sliding_window);
160+
const int kv_idx_min = std::max(0, diagonal - sliding_window);
180161
const int n_block_min = kv_idx_min / kBlockN;
181162
return make_tuple(n_block_min, n_block_max);
182163
} else {
@@ -300,31 +281,44 @@ struct FmhaBlock {
300281
}
301282

302283
// functions for tma load
284+
285+
// using StrideK = ...;
286+
287+
// using TMA_K = decltype(make_tma_copy(
288+
// GmemTiledCopy{}, // TMA_COPY
289+
// make_tensor(static_cast<InternalElementA const*>(nullptr),
290+
// repeat_like(StrideK{}, int32_t(0)), StrideK{}),
291+
// SmemLayoutK{}(_,_,_0{})));
292+
293+
// Tensor tensor_k = make_tensor(ptr_k, make_layout(make_shape(M,K,L),
294+
// args.stride_k)); auto tma_load_k = make_tma_copy(SM90_TMA_LOAD{},
295+
// gtensor_k, SmemLayoutK{}(_,_,_0{}));
296+
303297
// returns kv tma tile: (BLK_N, BLK_K, n) => (1@0, 1@1, 1@2)
304-
template <class TMA_K, class TMA_V>
305-
CUTE_HOST_DEVICE auto get_kv_tma_tile(TMA_K tma_k, TMA_V tma_v) const {
306-
// 1: make_gather_tma_tensor()
307-
// tma_tensor = (seq, dim, kv_head) => (1@0, 1@1, 1@2)
308-
// 2: partition into tiles
309-
// tma_tile = (BLK_N, BLK_K, n) => (1@0, 1@1, 1@2)
310-
311-
// (Q, D, (B, H))
312-
// Tensor mQ_qdl_p = tma_k.get_tma_tensor(select<0,2,3>(problem_shape));
313-
// Tensor mQ_qdl = domain_offset(make_coord(q_offs_0, _0{}, make_coord(_0{},
314-
// q_offs_2_1)), mQ_qdl_p); (BLK_N, BLK_K, m, k, (b)) Tensor gQ_qdl =
315-
// local_tile(mQ_qdl, TileShapeQK{}, make_coord(_, _, _), Step<_1, X,
316-
// _1>{});
317-
318-
// outside in caller part
319-
// (BLK_N, BLK_K, n) => (TMA,TMA_M,TMA_N, n)
320-
// auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
321-
// (TMA,TMA_M,TMA_N,REST_M,REST_N)
322-
// Tensor tAgA_x = cta_tma.partition_S(gA);
323-
// (TMA,TMA_M,TMA_N)
324-
// Tensor tAsA_x = cta_tma.partition_D(sA);
325-
326-
return;
327-
}
298+
// template <class TMA_K, class TMA_V>
299+
// CUTE_HOST_DEVICE auto get_kv_tma_tile(TMA_K tma_k, TMA_V tma_v) const {
300+
// 1: make_gather_tma_tensor()
301+
// tma_tensor = (seq, dim, kv_head) => (1@0, 1@1, 1@2)
302+
// 2: partition into tiles
303+
// tma_tile = (BLK_N, BLK_K, n) => (1@0, 1@1, 1@2)
304+
305+
// (Q, D, (B, H))
306+
// Tensor mQ_qdl_p = tma_k.get_tma_tensor(select<0,2,3>(problem_shape));
307+
// Tensor mQ_qdl = domain_offset(make_coord(q_offs_0, _0{}, make_coord(_0{},
308+
// q_offs_2_1)), mQ_qdl_p); (BLK_N, BLK_K, m, k, (b)) Tensor gQ_qdl =
309+
// local_tile(mQ_qdl, TileShapeQK{}, make_coord(_, _, _), Step<_1, X,
310+
// _1>{});
311+
312+
// outside in caller part
313+
// (BLK_N, BLK_K, n) => (TMA,TMA_M,TMA_N, n)
314+
// auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
315+
// (TMA,TMA_M,TMA_N,REST_M,REST_N)
316+
// Tensor tAgA_x = cta_tma.partition_S(gA);
317+
// (TMA,TMA_M,TMA_N)
318+
// Tensor tAsA_x = cta_tma.partition_D(sA);
319+
320+
// return;
321+
// }
328322
};
329323

330324
} // namespace llm

‎src/kernels/attention/common/tile_scheduler.cuh‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ namespace llm {
1010
using namespace cute;
1111
class SingleTileScheduler {
1212
public:
13+
// (m_block_idx, ((kv_head_idx, _0), batch_idx))
14+
using BlocKCoord = Coord<int, Coord<Coord<int, _0>, int>>;
15+
1316
// Device side kernel arguments
1417
struct Params {
1518
int batch_size = 0;
@@ -46,8 +49,7 @@ class SingleTileScheduler {
4649
Iterator() = default;
4750

4851
CUTE_DEVICE
49-
auto operator*() const {
50-
// (m_block_idx, ((kv_head_idx, _0), batch_idx))
52+
BlocKCoord operator*() const {
5153
return make_coord(
5254
(int)blockIdx.x,
5355
make_coord(make_coord((int)blockIdx.y, _0{}), (int)blockIdx.z));

‎src/kernels/attention/fmha_runner.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,6 @@ class FmhaRunner {
129129
.k_stride = k_stride,
130130
.v_stride = v_stride,
131131
.o_stride = o_stride,
132-
.sliding_window = params.sliding_window,
133132
},
134133
.mainloop =
135134
{

‎src/kernels/attention/kernel/builders/sm120_kernel_builder.inl‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,16 @@ struct KernelBuilder<cutlass::arch::Sm120,
3939
LOCAL,
4040
KV_USE_TMA,
4141
cute::enable_if_t<not cute::is_tuple_v<Element>>> {
42-
using Block =
43-
FmhaBlock<TileShape, Element, StrideQ, StrideK, StrideV, StrideO, LOCAL>;
42+
// TODO: support persistent kernels
43+
using TileScheduler = SingleTileScheduler;
44+
using BlocKCoord = TileScheduler::BlocKCoord;
45+
using Block = FmhaBlock<TileShape,
46+
BlocKCoord,
47+
Element,
48+
StrideQ,
49+
StrideK,
50+
StrideV,
51+
StrideO>;
4452

4553
using CollectiveMainloop = Sm120CollectiveFMhaWs<TileShape,
4654
Element,
@@ -55,9 +63,6 @@ struct KernelBuilder<cutlass::arch::Sm120,
5563
using CollectiveEpilogue =
5664
Sm120CollectiveEpilogue<TileShape, Element, EVEN_K>;
5765

58-
// TODO: support persistent kernels
59-
using TileScheduler = SingleTileScheduler;
60-
6166
using Kernel = Sm120KernelFmhaWs<ProblemShape,
6267
Block,
6368
CollectiveMainloop,

‎src/kernels/attention/kernel/sm120_kernel_fmha_ws.cuh‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -154,10 +154,15 @@ class Sm120KernelFmhaWs {
154154

155155
// process each block
156156
for (const auto blk_coord : scheduler) {
157-
// block coord: (batch_idx, m_block_idx, kv_head_idx)
158157
const Block block(params.block, blk_coord);
159-
mainloop.load(
160-
block, tidx, q_pipeline, q_state, kv_pipeline, kv_state, ss.mainloop);
158+
mainloop.load(params.mainloop,
159+
block,
160+
tidx,
161+
q_pipeline,
162+
q_state,
163+
kv_pipeline,
164+
kv_state,
165+
ss.mainloop);
161166
}
162167

163168
// prevent early exit of producer blocks in cluster
@@ -186,7 +191,6 @@ class Sm120KernelFmhaWs {
186191

187192
// process each block
188193
for (const auto blk_coord : scheduler) {
189-
// block coord: (batch_idx, m_block_idx, kv_head_idx)
190194
const Block block(params.block, blk_coord);
191195

192196
TiledMma tiled_mma;

0 commit comments

Comments
 (0)