Skip to content

[Feature] add support NoPE head dim 512 for GLM5Next model on SM90 #17

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zyongye merged 1 commit into from
Sep 15, 2026
Merged

[Feature] add support NoPE head dim 512 for GLM5Next model on SM90 #17

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions csrc/api/dense_decode.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,12 +182,20 @@ dense_attn_decode_interface(
params.stream = get_current_cuda_stream(q);

if (q_dtype == ScalarType::BFloat16) {
sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::bfloat16_t>(params);
if (head_size_k == 576) {
sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::bfloat16_t, 576>(params);
} else {
sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::bfloat16_t, 512>(params);
}
} else if (q_dtype == ScalarType::Half) {
#ifdef FLASH_MLA_DISABLE_FP16
STD_TORCH_CHECK(false, "FlashMLA is compiled with -DFLASH_MLA_DISABLE_FP16. Please remove this flag from your environment and re-compile FlashMLA.");
#else
sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::half_t>(params);
if (head_size_k == 576) {
sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::half_t, 576>(params);
} else {
sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::half_t, 512>(params);
}
#endif
} else {
STD_TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90");
Expand Down
1 change: 1 addition & 0 deletions csrc/kernels/sm90/decode/dense/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ static constexpr int BLOCK_SIZE_M = 64;
static constexpr int PAGE_BLOCK_SIZE = 64;

static constexpr int HEAD_DIM_K = 576;
static constexpr int HEAD_DIM_K_NOPE = 512;
static constexpr int HEAD_DIM_V = 512;

}
Expand Down
3 changes: 2 additions & 1 deletion csrc/kernels/sm90/decode/dense/instantiations/bf16.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

namespace sm90::decode::dense {

template void run_flash_splitkv_mla_kernel<cutlass::bfloat16_t>(DenseAttnDecodeParams &params);
template void run_flash_splitkv_mla_kernel<cutlass::bfloat16_t, 576>(DenseAttnDecodeParams &params);
template void run_flash_splitkv_mla_kernel<cutlass::bfloat16_t, 512>(DenseAttnDecodeParams &params);

}
3 changes: 2 additions & 1 deletion csrc/kernels/sm90/decode/dense/instantiations/fp16.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@
namespace sm90::decode::dense {

#ifndef FLASH_MLA_DISABLE_FP16
template void run_flash_splitkv_mla_kernel<cutlass::half_t>(DenseAttnDecodeParams &params);
template void run_flash_splitkv_mla_kernel<cutlass::half_t, 576>(DenseAttnDecodeParams &params);
template void run_flash_splitkv_mla_kernel<cutlass::half_t, 512>(DenseAttnDecodeParams &params);
#endif

}
74 changes: 46 additions & 28 deletions csrc/kernels/sm90/decode/dense/splitkv_mla.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,9 @@ __forceinline__ __device__ void warpgroup_cooperative_qkt_gemm(
QKT_GEMM_ONE_TILE(5);
QKT_GEMM_ONE_TILE(6);
QKT_GEMM_ONE_TILE(7);
QKT_GEMM_ONE_TILE(8);
if constexpr(T::REUSE_ROPE) {
QKT_GEMM_ONE_TILE(8);
}
QKT_GEMM_ONE_TILE(0);
QKT_GEMM_ONE_TILE(1);
QKT_GEMM_ONE_TILE(2);
Expand All @@ -248,7 +250,9 @@ __forceinline__ __device__ void warpgroup_cooperative_qkt_gemm(
QKT_GEMM_ONE_TILE(5);
QKT_GEMM_ONE_TILE(6);
QKT_GEMM_ONE_TILE(7);
QKT_GEMM_ONE_TILE(8);
if constexpr(T::REUSE_ROPE) {
QKT_GEMM_ONE_TILE(8);
}
cur_phase ^= 1;
}
}
Expand Down Expand Up @@ -705,7 +709,7 @@ __forceinline__ __device__ void launch_q_copy(
thr_tma.partition_S(my_tma_gQ),
thr_tma.partition_D(sQ)
);
barrier_Q->arrive_and_expect_tx(64*576*2);
barrier_Q->arrive_and_expect_tx(T::Q_BYTES);
}
}

Expand Down Expand Up @@ -754,8 +758,8 @@ __forceinline__ __device__ void wg0_subroutine(
Tensor<Engine11, Layout11> &rO0,
float rL[2],
int rRightBorderForQSeq[2],
TMABarrier barriers_K0[9],
TMABarrier barriers_K1[9],
TMABarrier barriers_K0[T::NUM_TILES_DIM_K],
TMABarrier barriers_K1[T::NUM_TILES_DIM_K],
bool &cur_phase_K0,
const TMAParams &tma_params,
const DenseAttnDecodeParams &params,
Expand Down Expand Up @@ -787,7 +791,7 @@ __forceinline__ __device__ void wg0_subroutine(

// Wait for rO0, launch TMA for the next V0L
cute::warpgroup_wait<0>();

// Wait for warpgroup 1, rescale P0, notify warpgroup 1
NamedBarrier::arrive_and_wait(T::NUM_THREADS, NamedBarriers::sScale1Ready);
if constexpr (!IS_BLK0_LAST && !IS_BLK1_LAST) {
Expand All @@ -798,7 +802,7 @@ __forceinline__ __device__ void wg0_subroutine(
save_rPb_to_sP<T>(rPb, sP0, idx_in_warpgroup);
cutlass::arch::fence_view_async_shared();
NamedBarrier::arrive(T::NUM_THREADS, NamedBarriers::sP0Ready);

// Wait for warpgroup 1, rescale O0, issue rO0 += rPb @ sV1L
if constexpr (!IS_BLK0_LAST) {
if constexpr (IS_BLK1_LAST) {
Expand All @@ -809,7 +813,7 @@ __forceinline__ __device__ void wg0_subroutine(
wg0_rescale_rO0(rO0, sScale1, rL, idx_in_warpgroup);
warpgroup_cooperative_pv_gemm_remoteP<T>(sP1, sV1L, rO0, idx_in_warpgroup);
}

// Issue P0 = Q @ K0^T
// Since TMAs for these 4 tiles are launched right after rO0 += rPb @ sV0L finishes, they should have already finished. Therefore, we issue the first 4 tiles to fill the pipeline.
if constexpr (!IS_BLK0_LAST && !IS_BLK1_LAST) {
Expand All @@ -821,12 +825,12 @@ __forceinline__ __device__ void wg0_subroutine(
cute::warpgroup_wait<4>();
launch_kv_tiles_copy_tma<0, 4>(tma_gK(_, _, nxt_block1_index), sK1, tma_params.tma_K, barriers_K1, idx_in_warpgroup);
}

// Issue P0 = Q @ K0^T
if constexpr (!IS_BLK0_LAST && !IS_BLK1_LAST) {
warpgroup_cooperative_qkt_gemm<T, 2>(sQ, sK0, rP0, rQ8, barriers_K0, cur_phase_K0, idx_in_warpgroup);
}

// Wait for P0 = Q @ K0^T
cute::warpgroup_wait<0>();
}
Expand Down Expand Up @@ -866,8 +870,8 @@ __forceinline__ __device__ void wg1_subroutine(
Tensor<Engine11, Layout11> &rO1,
float rL[2],
int rRightBorderForQSeq[2],
TMABarrier barriers_K0[9],
TMABarrier barriers_K1[9],
TMABarrier barriers_K0[T::NUM_TILES_DIM_K],
TMABarrier barriers_K1[T::NUM_TILES_DIM_K],
bool &cur_phase_K1,
const TMAParams &tma_params,
const DenseAttnDecodeParams &params,
Expand All @@ -882,7 +886,7 @@ __forceinline__ __device__ void wg1_subroutine(
int nxt_block1_index = GET_BLOCK_INDEX(block_idx+3);

Tensor rP1b = make_tensor<T::InputT>(Shape<Shape<_2, _2, _2>, _1, _4>{});

Tensor sV0R = get_half_V<T, 1>(sK0);
Tensor sV1R = get_half_V<T, 1>(sK1);

Expand All @@ -908,7 +912,7 @@ __forceinline__ __device__ void wg1_subroutine(
cutlass::arch::fence_view_async_shared();
}
}

// Wait for sP0, issue rO1 += sP0 @ sV0R, notify warpgroup 0
NamedBarrier::arrive_and_wait(T::NUM_THREADS, NamedBarriers::sP0Ready);
if constexpr (IS_BLK0_LAST) {
Expand All @@ -919,17 +923,17 @@ __forceinline__ __device__ void wg1_subroutine(
if constexpr (!IS_BLK0_LAST) {
NamedBarrier::arrive(T::NUM_THREADS, NamedBarriers::rO1sP0sV0RIssued);
}

// Wait for rO1 += rP1b @ sV1R, launch TMA for the next V1R
if constexpr (!IS_BLK0_LAST && !IS_BLK1_LAST && !IS_BLK2_LAST) {
cute::warpgroup_wait<1>();
launch_kv_tiles_copy_tma<4, 9>(tma_gK(_, _, nxt_block1_index), sK1, tma_params.tma_K, barriers_K1, idx_in_warpgroup);
launch_kv_tiles_copy_tma<4, T::NUM_TILES_DIM_K>(tma_gK(_, _, nxt_block1_index), sK1, tma_params.tma_K, barriers_K1, idx_in_warpgroup);
}

// Wait for rO1 += sP0 @ sV0R, launch TMA for the next V0R
if constexpr (!IS_BLK0_LAST && !IS_BLK1_LAST) {
cute::warpgroup_wait<0>();
launch_kv_tiles_copy_tma<4, 9>(tma_gK(_, _, nxt_block0_index), sK0, tma_params.tma_K, barriers_K0, idx_in_warpgroup);
launch_kv_tiles_copy_tma<4, T::NUM_TILES_DIM_K>(tma_gK(_, _, nxt_block0_index), sK0, tma_params.tma_K, barriers_K0, idx_in_warpgroup);
}

if constexpr (!IS_BLK0_LAST && !IS_BLK1_LAST && !IS_BLK2_LAST) {
Expand Down Expand Up @@ -982,8 +986,14 @@ flash_fwd_splitkv_mla_kernel(__grid_constant__ const DenseAttnDecodeParams param
Tensor sQ = make_tensor(make_smem_ptr(plan.smem_sQ.data()), (typename T::SmemLayoutQ){});
Tensor sK0 = make_tensor(make_smem_ptr(plan.smem_sK0.data()), (typename T::SmemLayoutK){});
Tensor sK1 = make_tensor(make_smem_ptr(plan.smem_sK1.data()), (typename T::SmemLayoutK){});
Tensor sP0 = make_tensor(make_smem_ptr(plan.smem_sP0.data()), (typename T::SmemLayoutP0){});
Tensor sP1 = flat_divide(sQ, Shape<Int<T::BLOCK_SIZE_M>, Int<T::PAGE_BLOCK_SIZE>>{})(_, _, _0{}, _8{}); // Overlap with sQ's 8-th tile
Tensor sP0 = make_tensor(make_smem_ptr(plan.smem_sP[0].data()), (typename T::SmemLayoutP0){});
Tensor sP1 = [&](){
if constexpr(T::REUSE_ROPE) {
return flat_divide(sQ, Shape<Int<T::BLOCK_SIZE_M>, Int<T::PAGE_BLOCK_SIZE>>{})(_, _, _0{}, _8{}); // Overlap with sQ's 8-th tile
} else {
return make_tensor(make_smem_ptr(plan.smem_sP[1].data()), (typename T::SmemLayoutP0){});
}
}();
Tensor sM = make_tensor(make_smem_ptr(plan.smem_sM.data()), make_shape(Int<T::BLOCK_SIZE_M>{}));
Tensor sL_reduction_wksp = make_tensor(make_smem_ptr(plan.sL_reduction_wksp.data()), make_shape(Int<2*T::BLOCK_SIZE_M>{}));
Tensor sScale0 = make_tensor(make_smem_ptr(plan.smem_sScale0.data()), make_shape(Int<T::BLOCK_SIZE_M>{}));
Expand All @@ -1007,7 +1017,7 @@ flash_fwd_splitkv_mla_kernel(__grid_constant__ const DenseAttnDecodeParams param
if (threadIdx.x == 0) {
barrier_Q->init(1);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < 9; ++i) {
for (int i = 0; i < T::NUM_TILES_DIM_K; ++i) {
barriers_K0[i].init(1);
barriers_K1[i].init(1);
}
Expand Down Expand Up @@ -1079,9 +1089,9 @@ flash_fwd_splitkv_mla_kernel(__grid_constant__ const DenseAttnDecodeParams param
>{});

// Copy K0 and K1
launch_kv_tiles_copy_tma<0, 9>(tma_gK(_, _, __ldg(block_table_ptr + start_block_idx)), sK0, tma_params.tma_K, barriers_K0, threadIdx.x);
launch_kv_tiles_copy_tma<0, T::NUM_TILES_DIM_K>(tma_gK(_, _, __ldg(block_table_ptr + start_block_idx)), sK0, tma_params.tma_K, barriers_K0, threadIdx.x);
if (start_block_idx+1 < end_block_idx) {
launch_kv_tiles_copy_tma<4, 9>(tma_gK(_, _, __ldg(block_table_ptr + start_block_idx+1)), sK1, tma_params.tma_K, barriers_K1, threadIdx.x);
launch_kv_tiles_copy_tma<4, T::NUM_TILES_DIM_K>(tma_gK(_, _, __ldg(block_table_ptr + start_block_idx+1)), sK1, tma_params.tma_K, barriers_K1, threadIdx.x);
launch_kv_tiles_copy_tma<0, 4>(tma_gK(_, _, __ldg(block_table_ptr + start_block_idx+1)), sK1, tma_params.tma_K, barriers_K1, threadIdx.x);
}

Expand All @@ -1100,7 +1110,9 @@ flash_fwd_splitkv_mla_kernel(__grid_constant__ const DenseAttnDecodeParams param
cur_phase_Q ^= 1;

Tensor rQ8 = make_tensor<InputT>(Shape<Shape<_2, _2, _2>, _1, _4>{});
retrieve_rP_from_sP<T>(rQ8, local_tile(sQ, Shape<_64, _64>{}, Coord<_0, _8>{}), idx_in_warpgroup);
if constexpr(T::REUSE_ROPE) {
retrieve_rP_from_sP<T>(rQ8, local_tile(sQ, Shape<_64, _64>{}, Coord<_0, _8>{}), idx_in_warpgroup);
}

if (warpgroup_idx == 0) {
// Warpgroup 0
Expand All @@ -1110,7 +1122,7 @@ flash_fwd_splitkv_mla_kernel(__grid_constant__ const DenseAttnDecodeParams param
// to a slow-down (or even register spilling, thanks to the great NVCC)
// Wait for K0
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < 9; ++i) {
for (int i = 0; i < T::NUM_TILES_DIM_K; ++i) {
if (idx_in_warpgroup == 0)
barriers_K0[i].arrive_and_expect_tx(64*64*2);
barriers_K0[i].wait(cur_phase_K0);
Expand Down Expand Up @@ -1272,12 +1284,18 @@ flash_fwd_splitkv_mla_kernel(__grid_constant__ const DenseAttnDecodeParams param
}


template<typename InputT>
template<typename InputT, int HEAD_DIM_K>
void run_flash_splitkv_mla_kernel(DenseAttnDecodeParams &params) {
FLASH_ASSERT(params.d == Config::HEAD_DIM_K);
if constexpr(HEAD_DIM_K > Config::HEAD_DIM_V) {
// with rope dim = 64
FLASH_ASSERT(params.d == Config::HEAD_DIM_K);
Comment on lines +1287 to +1291
} else {
// no rope
FLASH_ASSERT(params.d == Config::HEAD_DIM_K_NOPE);
}
FLASH_ASSERT(params.d_v == Config::HEAD_DIM_V);

using T = Traits<InputT>;
using T = Traits<InputT, HEAD_DIM_K>;
auto shape_Q = make_shape(params.q_seq_per_hk, params.d, params.h_k, params.b);
auto tma_Q = cute::make_tma_copy(
SM90_TMA_LOAD{},
Expand Down
2 changes: 1 addition & 1 deletion csrc/kernels/sm90/decode/dense/splitkv_mla.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

namespace sm90::decode::dense {

template<typename InputT>
template<typename InputT, int HEAD_DIM_K>
void run_flash_splitkv_mla_kernel(DenseAttnDecodeParams &params);

}
18 changes: 12 additions & 6 deletions csrc/kernels/sm90/decode/dense/traits.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,21 @@ namespace sm90::decode::dense {
using TMABarrier = cutlass::arch::ClusterTransactionBarrier;
using namespace cute;

template<typename InputT_>
template<typename InputT_, int HEAD_DIM_K_>
struct Traits {
using InputT = InputT_;

static constexpr int BLOCK_SIZE_M = Config::BLOCK_SIZE_M;
static constexpr int PAGE_BLOCK_SIZE = Config::PAGE_BLOCK_SIZE;
static constexpr int HEAD_DIM_K = Config::HEAD_DIM_K;
static constexpr int HEAD_DIM_K = HEAD_DIM_K_;
static constexpr int HEAD_DIM_V = Config::HEAD_DIM_V;

static constexpr int TILE_SIZE_K = 64;
static constexpr int REUSE_ROPE = (TILE_SIZE_K == PAGE_BLOCK_SIZE) \
&& (TILE_SIZE_K == (HEAD_DIM_K - HEAD_DIM_V));
static constexpr int Q_BYTES = BLOCK_SIZE_M * HEAD_DIM_K * sizeof(InputT);
static constexpr int NUM_TILES_DIM_K = HEAD_DIM_K / TILE_SIZE_K;

static constexpr int NUM_THREADS = 256;

static_assert(std::is_same_v<InputT, cutlass::bfloat16_t> || std::is_same_v<InputT, cutlass::half_t>);
Expand Down Expand Up @@ -74,13 +80,13 @@ struct Traits {
cute::array_aligned<InputT, cosize_v<SmemLayoutQ>> smem_sQ;
cute::array_aligned<InputT, cosize_v<SmemLayoutK>> smem_sK0;
cute::array_aligned<InputT, cosize_v<SmemLayoutK>> smem_sK1;
cute::array_aligned<InputT, cosize_v<SmemLayoutP0>> smem_sP0;
cute::array_aligned<InputT, cosize_v<SmemLayoutP0>> smem_sP[REUSE_ROPE ? 1 : 2];
cute::array_aligned<float, BLOCK_SIZE_M> smem_sM;
cute::array_aligned<float, 2*BLOCK_SIZE_M> sL_reduction_wksp;
cute::array_aligned<float, BLOCK_SIZE_M> smem_sScale0;
cute::array_aligned<float, BLOCK_SIZE_M> smem_sScale1;
TMABarrier barriers_K0[HEAD_DIM_K/64];
TMABarrier barriers_K1[HEAD_DIM_K/64];
TMABarrier barriers_K0[NUM_TILES_DIM_K];
TMABarrier barriers_K1[NUM_TILES_DIM_K];
TMABarrier barrier_Q;
};

Expand Down