Skip to content

Commit f6f726c

Browse files
zhoutianzi666zhoutianzi666
authored
clean code in sttantion (#3917)
1 parent 0d98982 commit f6f726c

4 files changed

Lines changed: 77 additions & 118 deletions

File tree

‎custom_ops/gpu_ops/append_attention.cu‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -317,7 +317,6 @@ void AppendAttentionKernel(
317317
qkv, // [token_num, num_heads, head_dim]
318318
seq_lens_decoder,
319319
seq_lens_encoder,
320-
batch_id_per_token,
321320
cu_seqlens_q,
322321
block_tables,
323322
rotary_embs,
@@ -344,7 +343,6 @@ void AppendAttentionKernel(
344343
qkv_out, // [token_num, num_heads, head_dim]
345344
seq_lens_decoder,
346345
seq_lens_encoder,
347-
batch_id_per_token,
348346
cu_seqlens_q,
349347
block_tables,
350348
rotary_embs,

‎custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_impl.cuh‎

Lines changed: 77 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,53 @@
1818
#include "mma_tensor_op.cuh"
1919
#include "utils.cuh"
2020

21+
22+
// Note(ZKK)
23+
// This function is very easy!
24+
// just make HeadDim data to be new HeadDim data!
25+
26+
template <typename T, int VecSize=8, int HEAD_DIM=128, int NUM_THREADS=32>
27+
__device__ __forceinline__ void apply_rope(
28+
const T* input,
29+
const float* cos_emb,
30+
const float* sin_emb,
31+
T* output,
32+
const int thread_id) {
33+
34+
using LoadT = AlignedVector<T, VecSize>;
35+
using LoadBiasT = AlignedVector<T, VecSize>;
36+
using LoadOutScaleT = AlignedVector<float, VecSize>;
37+
constexpr int HalfVecSize = VecSize / 2;
38+
using LoadEmbT = AlignedVector<float, HalfVecSize>;
39+
40+
LoadT src_vec;
41+
LoadBiasT out_vec;
42+
LoadEmbT cos_emb_vec;
43+
LoadEmbT sin_emb_vec;
44+
45+
#pragma unroll
46+
for (uint32_t head_bias = thread_id * VecSize; head_bias < HEAD_DIM; head_bias += NUM_THREADS * VecSize) {
47+
Load<T, VecSize>(&input[head_bias], &src_vec);
48+
const uint32_t emb_idx = head_bias / 2;
49+
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
50+
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
51+
#pragma unroll
52+
for (int i = 0; i < HalfVecSize; i++) {
53+
54+
float input_left = static_cast<float>(src_vec[2 * i]);
55+
float input_right = static_cast<float>(src_vec[2 * i + 1]);
56+
57+
const float cos_tmp = cos_emb_vec[i];
58+
const float sin_tmp = sin_emb_vec[i];
59+
out_vec[2 * i] =
60+
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
61+
out_vec[2 * i + 1] =
62+
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
63+
}
64+
Store<T, VecSize>(out_vec, &output[head_bias]);
65+
}
66+
}
67+
2168
template <typename T, int VecSize = 1>
2269
__global__ void append_decode_cache_T_rope_qk_norm_kernel(
2370
const T* __restrict__ quant_qkv, // [bsz, num_heads + 2 * kv_num_heads,
@@ -28,7 +75,7 @@ __global__ void append_decode_cache_T_rope_qk_norm_kernel(
2875
// head_size // 2]
2976
T* __restrict__ qkv_out,
3077
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
31-
const int* __restrict__ batch_id_per_token, // [num_tokens]
78+
3279
const int* __restrict__ cu_seqlens_q,
3380
const int* __restrict__ seq_lens, // [bsz]
3481
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -164,7 +211,7 @@ __global__ void append_decode_cache_T_rope_kernel(
164211
// head_size // 2]
165212
T* __restrict__ qkv_out,
166213
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
167-
const int* __restrict__ batch_id_per_token, // [num_tokens]
214+
168215
const int* __restrict__ cu_seqlens_q,
169216
const int* __restrict__ seq_lens, // [bsz]
170217
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -270,7 +317,7 @@ __global__ void append_decode_cache_T_rope_kernel(
270317
// head_size // 2]
271318
T* __restrict__ qkv_out,
272319
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
273-
const int* __restrict__ batch_id_per_token, // [num_tokens]
320+
274321
const int* __restrict__ cu_seqlens_q,
275322
const int* __restrict__ seq_lens, // [bsz]
276323
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -391,7 +438,6 @@ __global__ void append_decode_cache_T_neox_rope_kernel(
391438
// head_size // 2]
392439
T* __restrict__ qkv_out,
393440
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
394-
const int* __restrict__ batch_id_per_token, // [num_tokens]
395441
const int* __restrict__ cu_seqlens_q,
396442
const int* __restrict__ seq_lens, // [bsz]
397443
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -505,7 +551,6 @@ __global__ void append_decode_cache_T_neox_rope_kernel(
505551
// head_size // 2]
506552
T* __restrict__ qkv_out,
507553
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
508-
const int* __restrict__ batch_id_per_token, // [num_tokens]
509554
const int* __restrict__ cu_seqlens_q,
510555
const int* __restrict__ seq_lens, // [bsz]
511556
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -639,7 +684,6 @@ __global__ void append_decode_cache_int8_rope_kernel(
639684
// block_size, head_size // 2]
640685
T* __restrict__ qkv_out,
641686
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
642-
const int* __restrict__ batch_id_per_token, // [num_tokens]
643687
const int* __restrict__ cu_seqlens_q,
644688
const int* __restrict__ seq_lens, // [bsz]
645689
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -677,44 +721,18 @@ __global__ void append_decode_cache_int8_rope_kernel(
677721

678722
if (head_idx < num_heads) {
679723
// q
680-
using LoadT = AlignedVector<T, VecSize>;
681-
using LoadBiasT = AlignedVector<T, VecSize>;
682-
using LoadOutScaleT = AlignedVector<float, VecSize>;
683-
constexpr int HalfVecSize = VecSize / 2;
684-
using LoadEmbT = AlignedVector<float, HalfVecSize>;
724+
const T* qkv_now = quant_qkv + start_token_idx * hidden_size + head_idx * HeadDim;
725+
T* qkv_out_now = qkv_out + start_token_idx * hidden_size + head_idx * HeadDim;
726+
727+
uint32_t emb_offset = write_seq_id * half_head_size;
728+
emb_offset += rope_3d ? bid * max_seq_len * HeadDim : 0;
729+
apply_rope<T, VecSize, HeadDim, 32>(
730+
qkv_now,
731+
cos_emb + emb_offset,
732+
sin_emb + emb_offset,
733+
qkv_out_now,
734+
lane_id);
685735

686-
LoadT src_vec;
687-
LoadBiasT out_vec;
688-
LoadEmbT cos_emb_vec;
689-
LoadEmbT sin_emb_vec;
690-
const T* qkv_now = quant_qkv + start_token_idx * hidden_size;
691-
T* qkv_out_now = qkv_out + start_token_idx * hidden_size;
692-
#pragma unroll
693-
for (uint32_t head_bias = lane_id * VecSize; head_bias < HeadDim;
694-
head_bias += 32 * VecSize) {
695-
const int bias_idx = head_idx * HeadDim + head_bias;
696-
Load<T, VecSize>(&qkv_now[bias_idx], &src_vec);
697-
698-
// q rope
699-
const uint32_t emb_idx = write_seq_id * half_head_size + head_bias / 2;
700-
uint32_t new_emb_idx = rope_3d ? emb_idx + bid * max_seq_len * HeadDim : emb_idx;
701-
Load<float, HalfVecSize>(&cos_emb[new_emb_idx], &cos_emb_vec);
702-
Load<float, HalfVecSize>(&sin_emb[new_emb_idx], &sin_emb_vec);
703-
#pragma unroll
704-
for (int i = 0; i < HalfVecSize; i++) {
705-
// dequant + add_bias + rope
706-
float input_left = static_cast<float>(src_vec[2 * i]);
707-
float input_right = static_cast<float>(src_vec[2 * i + 1]);
708-
709-
const float cos_tmp = cos_emb_vec[i];
710-
const float sin_tmp = sin_emb_vec[i];
711-
out_vec[2 * i] =
712-
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
713-
out_vec[2 * i + 1] =
714-
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
715-
}
716-
Store<T, VecSize>(out_vec, &qkv_out_now[bias_idx]);
717-
}
718736
} else if (head_idx < num_heads + 2 * kv_num_heads) {
719737
// k
720738
constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t); // 16
@@ -889,7 +907,6 @@ __global__ void append_decode_cache_int8_rope_kernel(
889907
// block_size, head_size // 2]
890908
T* __restrict__ qkv_out,
891909
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
892-
const int* __restrict__ batch_id_per_token, // [num_tokens]
893910
const int* __restrict__ cu_seqlens_q,
894911
const int* __restrict__ seq_lens, // [bsz]
895912
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -1194,7 +1211,6 @@ __global__ void append_decode_cache_int8_neox_rope_kernel(
11941211
// block_size, head_size // 2]
11951212
T* __restrict__ qkv_out,
11961213
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
1197-
const int* __restrict__ batch_id_per_token, // [num_tokens]
11981214
const int* __restrict__ cu_seqlens_q,
11991215
const int* __restrict__ seq_lens, // [bsz]
12001216
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -1496,7 +1512,7 @@ __global__ void append_decode_cache_int8_neox_rope_kernel(
14961512
// block_size, head_size // 2]
14971513
T* __restrict__ qkv_out,
14981514
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
1499-
const int* __restrict__ batch_id_per_token, // [num_tokens]
1515+
15001516
const int* __restrict__ cu_seqlens_q,
15011517
const int* __restrict__ seq_lens, // [bsz]
15021518
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -1893,7 +1909,7 @@ __global__ void append_decode_cache_int4_rope_kernel(
18931909
// block_size, head_size // 2]
18941910
T* __restrict__ qkv_out,
18951911
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
1896-
const int* __restrict__ batch_id_per_token, // [num_tokens]
1912+
18971913
const int* __restrict__ cu_seqlens_q,
18981914
const int* __restrict__ seq_lens, // [bsz]
18991915
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -1934,44 +1950,18 @@ __global__ void append_decode_cache_int4_rope_kernel(
19341950

19351951
if (head_idx < num_heads) {
19361952
// q
1937-
using LoadT = AlignedVector<T, VecSize>;
1938-
using LoadBiasT = AlignedVector<T, VecSize>;
1939-
using LoadOutScaleT = AlignedVector<float, VecSize>;
1940-
constexpr int HalfVecSize = VecSize / 2;
1941-
using LoadEmbT = AlignedVector<float, HalfVecSize>;
1953+
const T* qkv_now = quant_qkv + start_token_idx * hidden_size + head_idx * HeadDim;
1954+
T* qkv_out_now = qkv_out + start_token_idx * hidden_size + head_idx * HeadDim;
1955+
1956+
uint32_t emb_offset = write_seq_id * half_head_size;
1957+
emb_offset += rope_3d ? bid * max_seq_len * HeadDim : 0;
1958+
apply_rope<T, VecSize, HeadDim, 32>(
1959+
qkv_now,
1960+
cos_emb + emb_offset,
1961+
sin_emb + emb_offset,
1962+
qkv_out_now,
1963+
lane_id);
19421964

1943-
LoadT src_vec;
1944-
LoadBiasT out_vec;
1945-
LoadEmbT cos_emb_vec;
1946-
LoadEmbT sin_emb_vec;
1947-
const T* qkv_now = quant_qkv + start_token_idx * hidden_size;
1948-
T* qkv_out_now = qkv_out + start_token_idx * hidden_size;
1949-
#pragma unroll
1950-
for (uint32_t head_bias = lane_id * VecSize; head_bias < HeadDim;
1951-
head_bias += 32 * VecSize) {
1952-
const int bias_idx = head_idx * HeadDim + head_bias;
1953-
Load<T, VecSize>(&qkv_now[bias_idx], &src_vec);
1954-
1955-
// q rope
1956-
const uint32_t emb_idx = write_seq_id * half_head_size + head_bias / 2;
1957-
uint32_t new_emb_idx = rope_3d ? emb_idx + bid * max_seq_len * HeadDim : emb_idx;
1958-
Load<float, HalfVecSize>(&cos_emb[new_emb_idx], &cos_emb_vec);
1959-
Load<float, HalfVecSize>(&sin_emb[new_emb_idx], &sin_emb_vec);
1960-
#pragma unroll
1961-
for (int i = 0; i < HalfVecSize; i++) {
1962-
// dequant + add_bias + rope
1963-
float input_left = static_cast<float>(src_vec[2 * i]);
1964-
float input_right = static_cast<float>(src_vec[2 * i + 1]);
1965-
1966-
const float cos_tmp = cos_emb_vec[i];
1967-
const float sin_tmp = sin_emb_vec[i];
1968-
out_vec[2 * i] =
1969-
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
1970-
out_vec[2 * i + 1] =
1971-
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
1972-
}
1973-
Store<T, VecSize>(out_vec, &qkv_out_now[bias_idx]);
1974-
}
19751965
} else if (head_idx < num_heads + 2 * kv_num_heads) {
19761966
// k
19771967
constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t); // 16
@@ -2191,7 +2181,7 @@ __global__ void append_decode_cache_int4_rope_kernel(
21912181
// block_size, head_size // 2]
21922182
T* __restrict__ qkv_out,
21932183
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
2194-
const int* __restrict__ batch_id_per_token, // [num_tokens]
2184+
21952185
const int* __restrict__ cu_seqlens_q,
21962186
const int* __restrict__ seq_lens, // [bsz]
21972187
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -2522,7 +2512,7 @@ __global__ void append_decode_cache_int4_neox_rope_kernel(
25222512
// block_size, head_size // 2]
25232513
T* __restrict__ qkv_out,
25242514
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
2525-
const int* __restrict__ batch_id_per_token, // [num_tokens]
2515+
25262516
const int* __restrict__ cu_seqlens_q,
25272517
const int* __restrict__ seq_lens, // [bsz]
25282518
const int* __restrict__ seq_lens_encoder, // [bsz]
@@ -2895,7 +2885,7 @@ __global__ void append_decode_cache_int4_neox_rope_kernel(
28952885
// block_size, head_size // 2]
28962886
T* __restrict__ qkv_out,
28972887
const int* __restrict__ block_tables, // [bsz, max_blocks_per_seq]
2898-
const int* __restrict__ batch_id_per_token, // [num_tokens]
2888+
28992889
const int* __restrict__ cu_seqlens_q,
29002890
const int* __restrict__ seq_lens, // [bsz]
29012891
const int* __restrict__ seq_lens_encoder, // [bsz]

0 commit comments

Comments
 (0)