ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 02:27:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use run-time checking (almost_invariant_seqlen) to replace pre-compiling setting (HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM) for switching grid-schedule

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-06-30 14:05:01 +00:00
parent 4667be2bb0
commit 53e69c1404
13 changed files with 341 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
example/ck_tile/18_hstu_attention/CMakeLists.txt
View File

@@ -12,10 +12,6 @@ set(EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS)
list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=3)
if (DEFINED ENV{ASSUME_HIGHLY_VARIED_SEQLEN})
list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -DHSTU_SCHED_BATCH_AS_FIRST_GRID_DIM=0)
endif()
if(GPU_TARGETS MATCHES "gfx95" AND NOT GPU_TARGETS MATCHES "gfx94" AND NOT GPU_TARGETS MATCHES "gfx90")
## disable slp-vectorize improve pipelines performance on gfx950
list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -DBUILD_HSTU_FOR_GFX95 -fno-slp-vectorize)

25
example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
View File

@@ -277,6 +277,18 @@ bool run_no_group_hstu_forward(const ck_tile::ArgParser& arg_parser, bool is_jag
phy_seqlen_kv = max_seqlen_kv;
};
int min_seqlen_q = std::numeric_limits<int>::max();
int min_seqlen_kv = std::numeric_limits<int>::max();
if(is_jagged)
{
for(int i = 0; i < num_batch; i++)
{
min_seqlen_q = min(min_seqlen_q, seq_offsets_q[i + 1] - seq_offsets_q[i]);
min_seqlen_kv = min(min_seqlen_kv, seq_offsets_kv[i + 1] - seq_offsets_kv[i]);
};
};
long total_flops = 0;
// estimate the total flops occurred, ignoring the scaling and SiLu
@@ -379,6 +391,8 @@ bool run_no_group_hstu_forward(const ck_tile::ArgParser& arg_parser, bool is_jag
params.seq_kv_offsets_ptr = seq_offsets_kv_dev.GetDeviceBuffer();
params.max_seqlen_q = max_seqlen_q;
params.max_seqlen_kv = max_seqlen_kv;
params.min_seqlen_q = min_seqlen_q;
params.min_seqlen_kv = min_seqlen_kv;
params.q_ptr = q_dev.GetDeviceBuffer();
params.k_ptr = k_dev.GetDeviceBuffer();
params.v_ptr = v_dev.GetDeviceBuffer();
@@ -798,6 +812,15 @@ bool run_group_hstu_forward(const ck_tile::ArgParser& arg_parser, int num_group)
}
};
int min_seqlen_q = std::numeric_limits<int>::max();
int min_seqlen_kv = std::numeric_limits<int>::max();
for(int i = 0; i < num_batch; i++)
{
min_seqlen_q = min(min_seqlen_q, seq_offsets_q[i + 1] - seq_offsets_q[i]);
min_seqlen_kv = min(min_seqlen_kv, seq_offsets_kv[i + 1] - seq_offsets_kv[i]);
};
long total_flops = 0;
// estimate the total flops occurred, ignoring the scaling and SILu
@@ -900,6 +923,8 @@ bool run_group_hstu_forward(const ck_tile::ArgParser& arg_parser, int num_group)
params.seq_kv_offsets_ptr = seq_offsets_kv_dev.GetDeviceBuffer();
params.max_seqlen_q = max_max_seqlen_q;
params.max_seqlen_kv = max_max_seqlen_kv;
params.min_seqlen_q = min_seqlen_q;
params.min_seqlen_kv = min_seqlen_kv;
params.q_ptr = q_dev.GetDeviceBuffer();
params.k_ptr = k_dev.GetDeviceBuffer();
params.v_ptr = v_dev.GetDeviceBuffer();

10
example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_dispatch.hpp
View File

@@ -176,9 +176,13 @@ struct batched_forward_dispatch
param.philox_offset);
}();
bool has_minfull_attn_seqlen = (param.min_full_attn_seqlen > 0);
dim3 kGridSize = HstuKernel::GridSize(
param.num_batch, param.num_head, param.seqlen_q, param.hdim_v, has_minfull_attn_seqlen);
bool has_minfull_attn_seqlen = (param.min_full_attn_seqlen > 0);
dim3 kGridSize = HstuKernel::GridSize(param.num_batch,
param.num_head,
param.seqlen_q,
param.hdim_v,
true /* almost_invariant_seqlen */,
has_minfull_attn_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;

6
example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_splitkv_dispatch.hpp
View File

@@ -316,6 +316,7 @@ struct batched_forward_splitkv_dispatch
param.seqlen_q,
param.hdim_v,
ws.num_splits,
true, // almost_invariant_seqlen
has_minfull_attn_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;
@@ -347,8 +348,9 @@ struct batched_forward_splitkv_dispatch
param.hdim_v);
}();
dim3 kGridSize = HstuKernel::GridSize(param.num_batch, param.num_head, param.seqlen_q);
dim3 kBlockSize = HstuKernel::BlockSize();
dim3 kGridSize = HstuKernel::GridSize(
param.num_batch, param.num_head, param.seqlen_q, true /* almost_invariant_seqlen */);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;
(void)ck_tile::launch_kernel(

151
example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp
View File

@@ -14,10 +14,6 @@
#include "hstu_block_masking.hpp"
#include "hstu_attention_kernel_util.hpp"
#ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1
#endif
// S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q]
// S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1]
// S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k]
@@ -101,6 +97,8 @@ struct HstuAttentionFwdKernel
float scale_s; // scaling value exerted on the immediate Q@K result
float scale_p; // scaling value exerted on the SiLU result
bool almost_invariant_seqlen; // should always be true for batched mode
ck_tile::index_t contextual_seqlen;
ck_tile::index_t window_size;
ck_tile::index_t min_full_attn_seqlen;
@@ -138,6 +136,8 @@ struct HstuAttentionFwdKernel
float scale_s; // scaling value exerted on the immediate Q@K result
float scale_p; // scaling value exerted on the SiLU result
bool almost_invariant_seqlen;
ck_tile::index_t contextual_seqlen;
ck_tile::index_t window_size;
ck_tile::index_t min_full_attn_seqlen;
@@ -177,6 +177,8 @@ struct HstuAttentionFwdKernel
float scale_s; // scaling value exerted on the immediate Q@K result
float scale_p; // scaling value exerted on the SiLU result
bool almost_invariant_seqlen;
int32_t contextual_seqlen; // to be set by the per-group contextual_seqlen
int32_t window_size; // to be set by the per-group window_size
int32_t min_full_attn_seqlen; // to be set by the per-group min_full_attn_seqlen
@@ -356,6 +358,7 @@ struct HstuAttentionFwdKernel
num_head,
scale_s,
attn_scale ? attn_scale : 1.0f / static_cast<float>(seqlen_q),
true, // almost_invariant_seqlen
contextual_seqlen,
window_size,
min_full_attn_seqlen}, // args for common karg
@@ -402,6 +405,7 @@ struct HstuAttentionFwdKernel
ck_tile::index_t num_head,
float scale_s,
float attn_scale,
bool almost_invariant_seqlen,
ck_tile::index_t seq_stride_q,
ck_tile::index_t seq_stride_k,
ck_tile::index_t seq_stride_v,
@@ -445,6 +449,7 @@ struct HstuAttentionFwdKernel
num_head,
scale_s,
attn_scale ? attn_scale : 1.0f / static_cast<float>(max_seqlen_q),
almost_invariant_seqlen,
contextual_seqlen,
window_size,
min_full_attn_seqlen}, // args for common karg
@@ -493,6 +498,7 @@ struct HstuAttentionFwdKernel
ck_tile::index_t hdim_v,
ck_tile::index_t num_head,
float scale_s,
bool almost_invariant_seqlen,
ck_tile::index_t seq_stride_q,
ck_tile::index_t seq_stride_k,
ck_tile::index_t seq_stride_v,
@@ -534,9 +540,10 @@ struct HstuAttentionFwdKernel
num_head,
scale_s,
1.0f, // to be set according to the per-group attn_scale and max_seqlen
0, // to be set by the per-group contextual_seqlen
0, // to be set by the per-group window_size
0, // to be set by the per-group min_full_attn_seqlen
almost_invariant_seqlen,
0, // to be set by the per-group contextual_seqlen
0, // to be set by the per-group window_size
0, // to be set by the per-group min_full_attn_seqlen
reinterpret_cast<const int32_t*>(group_max_seqlen_q_ptr),
reinterpret_cast<const int32_t*>(group_contextual_seqlen_ptr),
reinterpret_cast<const int32_t*>(group_window_size_ptr),
@@ -571,6 +578,7 @@ struct HstuAttentionFwdKernel
ck_tile::index_t nhead_,
ck_tile::index_t seqlen_,
ck_tile::index_t hdim_v_,
bool almost_invariant_seqlen,
bool has_minfull_attn_seqlen = false)
{
// The Q sequence [0, seqlen) will be split to two parts for allocating workgroups:
@@ -589,27 +597,33 @@ struct HstuAttentionFwdKernel
num_tile_in_seqlen += 1;
};
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
if(almost_invariant_seqlen)
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_,
nhead_,
num_tile_in_seqlen *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1));
#else
return dim3(num_tile_in_seqlen *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1),
nhead_,
batch_size_);
#endif
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
return dim3(batch_size_,
nhead_,
num_tile_in_seqlen *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1));
}
else
{
return dim3(batch_size_, nhead_, num_tile_in_seqlen);
}
}
else
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_, nhead_, num_tile_in_seqlen);
#else
return dim3(num_tile_in_seqlen, nhead_, batch_size_);
#endif
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
return dim3(num_tile_in_seqlen *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1),
nhead_,
batch_size_);
}
else
{
return dim3(num_tile_in_seqlen, nhead_, batch_size_);
}
}
}
@@ -620,52 +634,65 @@ struct HstuAttentionFwdKernel
const index_t num_tile_n1 =
ck_tile::integer_divide_ceil(kargs.hdim_v, HstuAttentionPipeline::kN1);
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#else
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
#endif
if(kargs.almost_invariant_seqlen)
{
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
const auto f = [](index_t dividend, index_t divisor) {
index_t quotient = dividend / divisor;
index_t modulus = dividend - quotient * divisor;
return ck_tile::make_tuple(quotient, modulus);
};
const auto f = [](index_t dividend, index_t divisor) {
index_t quotient = dividend / divisor;
index_t modulus = dividend - quotient * divisor;
return ck_tile::make_tuple(quotient, modulus);
};
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1);
i_tile_m = gridDim.z / num_tile_n1 - 1 - i_tile_m;
#else
const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1);
#endif
auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1);
i_tile_m = gridDim.z / num_tile_n1 - 1 - i_tile_m;
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
}
else
{
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
const auto f = [](index_t dividend, index_t divisor) {
index_t quotient = dividend / divisor;
index_t modulus = dividend - quotient * divisor;
return ck_tile::make_tuple(quotient, modulus);
};
const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1);
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
}
}
else
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#else
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
#endif
if(kargs.almost_invariant_seqlen)
{
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
index_t i_tile_m = i_block;
i_tile_m = gridDim.z - 1 - i_tile_m;
#else
const index_t i_tile_m = i_block;
#endif
const index_t i_tile_n = 0;
index_t i_tile_m = i_block;
i_tile_m = gridDim.z - 1 - i_tile_m;
const index_t i_tile_n = 0;
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
}
else
{
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
const index_t i_tile_m = i_block;
const index_t i_tile_n = 0;
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
}
}
}

66
example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_combine_kernel.hpp
View File

@@ -14,10 +14,6 @@
#include "hstu_block_masking.hpp"
#ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1
#endif
// S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q]
// S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1]
// S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k]
@@ -71,6 +67,8 @@ struct HstuAttentionFwdSplitKVCombineKernel
ck_tile::index_t num_head;
ck_tile::index_t num_splits;
ck_tile::index_t hdim_v;
bool almost_invariant_seqlen; // should always be true for batched mode
};
struct HstuAttentionJaggedCombineBaseKargs
@@ -87,6 +85,8 @@ struct HstuAttentionFwdSplitKVCombineKernel
ck_tile::index_t hdim_v;
ck_tile::index_t seqlen_q;
bool almost_invariant_seqlen;
};
struct HstuAttentionCombineSoftmaxKargs
@@ -160,9 +160,10 @@ struct HstuAttentionFwdSplitKVCombineKernel
seqlen_q,
num_head,
num_splits,
hdim_v},
{}, // place holder for softmax
{}, // place holder for LSE
hdim_v,
true}, // almost_invariant_seqlen
{}, // place holder for softmax
{}, // place holder for LSE
};
if constexpr(kUseSoftmax)
@@ -193,7 +194,8 @@ struct HstuAttentionFwdSplitKVCombineKernel
const void* seq_q_offsets_ptr,
ck_tile::index_t num_head,
ck_tile::index_t num_splits, // number of splitted seqlen_kv
ck_tile::index_t hdim_v)
ck_tile::index_t hdim_v,
bool almost_invariant_seqlen)
{
Kargs kargs{
{o_acc_ptr,
@@ -204,7 +206,8 @@ struct HstuAttentionFwdSplitKVCombineKernel
num_head,
num_splits,
hdim_v,
0 /* seqlen_q will be updated later*/},
0, /* seqlen_q will be updated later*/
almost_invariant_seqlen},
{}, // place holder for softmax
{}, // place holder for LSE
};
@@ -223,33 +226,40 @@ struct HstuAttentionFwdSplitKVCombineKernel
return kargs;
}
CK_TILE_HOST static constexpr auto
GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_)
CK_TILE_HOST static constexpr auto GridSize(ck_tile::index_t batch_size_,
ck_tile::index_t nhead_,
ck_tile::index_t seqlen_,
bool almost_invariant_seqlen)
{
ck_tile::index_t num_tile_in_seqlen =
ck_tile::integer_divide_ceil(seqlen_, HstuAttentionPipeline::kM);
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_, nhead_, num_tile_in_seqlen);
#else
return dim3(num_tile_in_seqlen, nhead_, batch_size_);
#endif
if(almost_invariant_seqlen)
{
return dim3(batch_size_, nhead_, num_tile_in_seqlen);
}
else
{
return dim3(num_tile_in_seqlen, nhead_, batch_size_);
}
}
CK_TILE_DEVICE static constexpr auto GetTileIndex(const Kargs& kargs)
{
ignore = kargs;
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_tile_m = blockIdx.z;
#else
const index_t i_tile_m = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
#endif
return ck_tile::make_tuple(i_tile_m, i_nhead, i_batch);
if(kargs.almost_invariant_seqlen)
{
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_tile_m = blockIdx.z;
return ck_tile::make_tuple(i_tile_m, i_nhead, i_batch);
}
else
{
const index_t i_tile_m = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
return ck_tile::make_tuple(i_tile_m, i_nhead, i_batch);
}
}
CK_TILE_HOST static constexpr auto BlockSize()

143
example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_kernel.hpp
View File

@@ -14,10 +14,6 @@
#include "hstu_block_masking.hpp"
#include "hstu_attention_kernel_util.hpp"
#ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1
#endif
// S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q]
// S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1]
// S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k]
@@ -99,6 +95,8 @@ struct HstuAttentionFwdSplitKVKernel
float scale_s; // scaling value exerted on the immediate Q@K result
float scale_p; // scaling value exerted on the SiLU result
bool almost_invariant_seqlen; // should always be true for batched mode
ck_tile::index_t contextual_seqlen;
ck_tile::index_t window_size;
ck_tile::index_t min_full_attn_seqlen;
@@ -135,6 +133,8 @@ struct HstuAttentionFwdSplitKVKernel
float scale_s; // scaling value exerted on the immediate Q@K result
float scale_p; // scaling value exerted on the SiLU result
bool almost_invariant_seqlen;
ck_tile::index_t contextual_seqlen;
ck_tile::index_t window_size;
ck_tile::index_t min_full_attn_seqlen;
@@ -173,6 +173,8 @@ struct HstuAttentionFwdSplitKVKernel
float scale_s; // scaling value exerted on the immediate Q@K result
float scale_p; // scaling value exerted on the SiLU result
bool almost_invariant_seqlen;
int32_t contextual_seqlen; // to be set by the per-group contextual_seqlen
int32_t window_size; // to be set by the per-group window_size
int32_t min_full_attn_seqlen; // to be set by the per-group min_full_attn_seqlen
@@ -334,6 +336,7 @@ struct HstuAttentionFwdSplitKVKernel
num_head,
scale_s,
attn_scale ? attn_scale : 1.0f / static_cast<float>(seqlen_q),
true, // almost_invariant_seqlen
contextual_seqlen,
window_size,
min_full_attn_seqlen}, // args for common karg
@@ -377,6 +380,7 @@ struct HstuAttentionFwdSplitKVKernel
ck_tile::index_t num_head,
float scale_s,
float attn_scale,
bool almost_invariant_seqlen,
ck_tile::index_t seq_stride_q,
ck_tile::index_t seq_stride_k,
ck_tile::index_t seq_stride_v,
@@ -415,6 +419,7 @@ struct HstuAttentionFwdSplitKVKernel
num_head,
scale_s,
attn_scale ? attn_scale : 1.0f / static_cast<float>(max_seqlen_q),
almost_invariant_seqlen,
contextual_seqlen,
window_size,
min_full_attn_seqlen}, // args for common karg
@@ -461,6 +466,7 @@ struct HstuAttentionFwdSplitKVKernel
ck_tile::index_t hdim_v,
ck_tile::index_t num_head,
float scale_s,
bool almost_invariant_seqlen,
ck_tile::index_t seq_stride_q,
ck_tile::index_t seq_stride_k,
ck_tile::index_t seq_stride_v,
@@ -497,9 +503,10 @@ struct HstuAttentionFwdSplitKVKernel
num_head,
scale_s,
1.0f, // to be set according to the per-group attn_scale and max_seqlen
0, // to be set by the per-group contextual_seqlen
0, // to be set by the per-group window_size
0, // to be set by the per-group min_full_attn_seqlen
almost_invariant_seqlen,
0, // to be set by the per-group contextual_seqlen
0, // to be set by the per-group window_size
0, // to be set by the per-group min_full_attn_seqlen
reinterpret_cast<const int32_t*>(group_max_seqlen_q_ptr),
reinterpret_cast<const int32_t*>(group_contextual_seqlen_ptr),
reinterpret_cast<const int32_t*>(group_window_size_ptr),
@@ -532,6 +539,7 @@ struct HstuAttentionFwdSplitKVKernel
ck_tile::index_t seqlen_,
ck_tile::index_t hdim_v_,
ck_tile::index_t num_splits,
bool almost_invariant_seqlen,
bool has_minfull_attn_seqlen = false)
{
// The Q sequence [0, seqlen) will be split to two parts for allocating workgroups:
@@ -550,27 +558,33 @@ struct HstuAttentionFwdSplitKVKernel
num_tile_in_seqlen += 1;
};
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
if(almost_invariant_seqlen)
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_,
nhead_,
num_tile_in_seqlen * num_splits *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1));
#else
return dim3(num_tile_in_seqlen * num_splits *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1),
nhead_,
batch_size_);
#endif
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
return dim3(batch_size_,
nhead_,
num_tile_in_seqlen * num_splits *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1));
}
else
{
return dim3(batch_size_, nhead_, num_tile_in_seqlen * num_splits);
}
}
else
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_, nhead_, num_tile_in_seqlen * num_splits);
#else
return dim3(num_tile_in_seqlen * num_splits, nhead_, batch_size_);
#endif
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
return dim3(num_tile_in_seqlen * num_splits *
ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1),
nhead_,
batch_size_);
}
else
{
return dim3(num_tile_in_seqlen * num_splits, nhead_, batch_size_);
}
}
}
@@ -587,51 +601,58 @@ struct HstuAttentionFwdSplitKVKernel
const index_t num_tile_n1 =
ck_tile::integer_divide_ceil(kargs.hdim_v, HstuAttentionPipeline::kN1);
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#else
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
#endif
if(kargs.almost_invariant_seqlen)
{
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1);
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1);
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
i_tile_m = gridDim.z / num_tile_n1 / kargs.num_splits - 1 - i_tile_m;
#else
auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1);
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
#endif
i_tile_m = gridDim.z / num_tile_n1 / kargs.num_splits - 1 - i_tile_m;
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split);
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split);
}
else
{
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1);
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split);
}
}
else
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#else
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
#endif
if(kargs.almost_invariant_seqlen)
{
const index_t i_batch = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_block = blockIdx.z;
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
index_t i_tile_m_i_split = i_block;
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
i_tile_m = gridDim.z / kargs.num_splits - 1 - i_tile_m;
#else
index_t i_tile_m_i_split = i_block;
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
#endif
const index_t i_tile_n = 0;
index_t i_tile_m_i_split = i_block;
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
i_tile_m = gridDim.z / kargs.num_splits - 1 - i_tile_m;
const index_t i_tile_n = 0;
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split);
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split);
}
else
{
const index_t i_block = blockIdx.x;
const index_t i_nhead = blockIdx.y;
const index_t i_batch = blockIdx.z;
index_t i_tile_m_i_split = i_block;
auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits);
const index_t i_tile_n = 0;
return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split);
}
}
}

10
example/ck_tile/18_hstu_attention/hstu_attention_group_forward_dispatch.hpp
View File

@@ -124,6 +124,8 @@ struct group_forward_dispatch
template <typename HstuKernel>
static void RunWithKernel(HstuAttentionGroupFwdParams& param, hipStream_t stream)
{
bool almost_invariant_seqlen = is_almost_invariant_seqlen(param);
const auto kargs = [&] {
return HstuKernel::MakeKargs(param.q_ptr,
param.k_ptr,
@@ -144,6 +146,7 @@ struct group_forward_dispatch
param.hdim_v,
param.num_head,
param.scale_s,
almost_invariant_seqlen,
param.seq_stride_q,
param.seq_stride_k,
param.seq_stride_v,
@@ -162,8 +165,11 @@ struct group_forward_dispatch
param.philox_offset);
}();
dim3 kGridSize =
HstuKernel::GridSize(param.num_batch, param.num_head, param.max_seqlen_q, param.hdim_v);
dim3 kGridSize = HstuKernel::GridSize(param.num_batch,
param.num_head,
param.max_seqlen_q,
param.hdim_v,
almost_invariant_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;

21
example/ck_tile/18_hstu_attention/hstu_attention_group_forward_splitkv_dispatch.hpp
View File

@@ -265,6 +265,8 @@ struct group_forward_splitkv_dispatch
HIP_CHECK_ERROR(hipMallocAsync(&ws.lse_acc_ptr, workspace_bytes, stream));
}
bool almost_invariant_seqlen = is_almost_invariant_seqlen(param);
const auto kargs = [&] {
return HstuKernel::MakeKargs(param.q_ptr,
param.k_ptr,
@@ -286,6 +288,7 @@ struct group_forward_splitkv_dispatch
param.hdim_v,
param.num_head,
param.scale_s,
almost_invariant_seqlen,
param.seq_stride_q,
param.seq_stride_k,
param.seq_stride_v,
@@ -300,8 +303,12 @@ struct group_forward_splitkv_dispatch
param.philox_offset);
}();
dim3 kGridSize = HstuKernel::GridSize(
param.num_batch, param.num_head, param.max_seqlen_q, param.hdim_v, ws.num_splits);
dim3 kGridSize = HstuKernel::GridSize(param.num_batch,
param.num_head,
param.max_seqlen_q,
param.hdim_v,
ws.num_splits,
almost_invariant_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;
@@ -315,6 +322,8 @@ struct group_forward_splitkv_dispatch
SplitkvWorkspace& ws,
hipStream_t stream)
{
bool almost_invariant_seqlen = is_almost_invariant_seqlen_q(param);
const auto kargs = [&] {
return HstuKernel::MakeKargs(ws.o_acc_ptr,
ws.lse_acc_ptr,
@@ -327,11 +336,13 @@ struct group_forward_splitkv_dispatch
param.seq_q_offsets_ptr,
param.num_head,
ws.num_splits,
param.hdim_v);
param.hdim_v,
almost_invariant_seqlen);
}();
dim3 kGridSize = HstuKernel::GridSize(param.num_batch, param.num_head, param.max_seqlen_q);
dim3 kBlockSize = HstuKernel::BlockSize();
dim3 kGridSize = HstuKernel::GridSize(
param.num_batch, param.num_head, param.max_seqlen_q, almost_invariant_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;
(void)ck_tile::launch_kernel(

52
example/ck_tile/18_hstu_attention/hstu_attention_host_util.hpp
View File

@@ -9,6 +9,7 @@
#include <stdexcept>
#include <ck_tile/host/hip_check_error.hpp>
#include "hstu_attention_params.hpp"
#define HSTU_CHECK(COND, ERR) \
if(!(COND)) \
@@ -30,3 +31,54 @@ static inline int get_number_of_cu()
return props.multiProcessorCount;
}
static inline bool is_almost_invariant_seqlen(HstuAttentionNoGroupFwdParams& param)
{
float threshold = 0.7f;
if(param.is_jagged)
{
bool res = (static_cast<float>(param.min_seqlen_q) / param.max_seqlen_q) > threshold;
if(param.is_cross_attention)
res = res &&
((static_cast<float>(param.min_seqlen_kv) / param.max_seqlen_kv) > threshold);
return res;
}
else
return true;
};
static inline bool is_almost_invariant_seqlen(HstuAttentionGroupFwdParams& param)
{
float threshold = 0.7f;
bool res = (static_cast<float>(param.min_seqlen_q) / param.max_seqlen_q) > threshold;
if(param.is_cross_attention)
res = res && ((static_cast<float>(param.min_seqlen_kv) / param.max_seqlen_kv) > threshold);
return res;
};
static inline bool is_almost_invariant_seqlen_q(HstuAttentionNoGroupFwdParams& param)
{
float threshold = 0.7f;
if(param.is_jagged)
{
bool res = (static_cast<float>(param.min_seqlen_q) / param.max_seqlen_q) > threshold;
return res;
}
else
return true;
};
static inline bool is_almost_invariant_seqlen_q(HstuAttentionGroupFwdParams& param)
{
float threshold = 0.7f;
bool res = (static_cast<float>(param.min_seqlen_q) / param.max_seqlen_q) > threshold;
return res;
};

5
example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_dispatch.hpp
View File

@@ -23,6 +23,7 @@
#include "hstu_attention_epilogue.hpp"
#include "hstu_attention_jagged_forward_splitkv_dispatch.hpp"
#include "hstu_attention_host_util.hpp"
template <typename InOutDataType,
bool kUseCausal,
@@ -124,6 +125,8 @@ struct jagged_forward_dispatch
template <typename HstuKernel>
static void RunWithKernel(HstuAttentionNoGroupFwdParams& param, hipStream_t stream)
{
bool almost_invariant_seqlen = is_almost_invariant_seqlen(param);
const auto kargs = [&] {
return HstuKernel::MakeKargs(param.q_ptr,
param.k_ptr,
@@ -140,6 +143,7 @@ struct jagged_forward_dispatch
param.num_head,
param.scale_s,
param.attn_scale,
almost_invariant_seqlen,
param.seq_stride_q,
param.seq_stride_k,
param.seq_stride_v,
@@ -166,6 +170,7 @@ struct jagged_forward_dispatch
param.num_head,
param.max_seqlen_q,
param.hdim_v,
almost_invariant_seqlen,
has_minfull_attn_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;

14
example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_splitkv_dispatch.hpp
View File

@@ -264,6 +264,8 @@ struct jagged_forward_splitkv_dispatch
HIP_CHECK_ERROR(hipMallocAsync(&ws.lse_acc_ptr, workspace_bytes, stream));
}
bool almost_invariant_seqlen = is_almost_invariant_seqlen(param);
const auto kargs = [&] {
return HstuKernel::MakeKargs(param.q_ptr,
param.k_ptr,
@@ -281,6 +283,7 @@ struct jagged_forward_splitkv_dispatch
param.num_head,
param.scale_s,
param.attn_scale,
almost_invariant_seqlen,
param.seq_stride_q,
param.seq_stride_k,
param.seq_stride_v,
@@ -304,6 +307,7 @@ struct jagged_forward_splitkv_dispatch
param.max_seqlen_q,
param.hdim_v,
ws.num_splits,
almost_invariant_seqlen,
has_minfull_attn_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;
@@ -318,6 +322,8 @@ struct jagged_forward_splitkv_dispatch
SplitkvWorkspace& ws,
hipStream_t stream)
{
bool almost_invariant_seqlen = is_almost_invariant_seqlen_q(param);
const auto kargs = [&] {
return HstuKernel::MakeKargs(ws.o_acc_ptr,
ws.lse_acc_ptr,
@@ -330,11 +336,13 @@ struct jagged_forward_splitkv_dispatch
param.seq_q_offsets_ptr,
param.num_head,
ws.num_splits,
param.hdim_v);
param.hdim_v,
almost_invariant_seqlen);
}();
dim3 kGridSize = HstuKernel::GridSize(param.num_batch, param.num_head, param.max_seqlen_q);
dim3 kBlockSize = HstuKernel::BlockSize();
dim3 kGridSize = HstuKernel::GridSize(
param.num_batch, param.num_head, param.max_seqlen_q, almost_invariant_seqlen);
dim3 kBlockSize = HstuKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu;
(void)ck_tile::launch_kernel(

4
example/ck_tile/18_hstu_attention/hstu_attention_params.hpp
View File

@@ -25,6 +25,8 @@ struct HstuAttentionNoGroupFwdParams
const void* seq_kv_offsets_ptr; // jagged mode only
ck_tile::index_t max_seqlen_q; // jagged mode only
ck_tile::index_t max_seqlen_kv; // jagged mode only
ck_tile::index_t min_seqlen_q; // jagged mode only
ck_tile::index_t min_seqlen_kv; // jagged mode only
const void* q_ptr;
const void* k_ptr;
@@ -90,6 +92,8 @@ struct HstuAttentionGroupFwdParams
const void* seq_kv_offsets_ptr;
ck_tile::index_t max_seqlen_q; // the maximum of all the groups' max_seqlen_q
ck_tile::index_t max_seqlen_kv; // the maximum of all the groups' max_seqlen_kv
ck_tile::index_t min_seqlen_q; // jagged mode only
ck_tile::index_t min_seqlen_kv; // jagged mode only
const void* q_ptr;
const void* k_ptr;