minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] Fix issues when running llama3.2-11B vision model with image tasks (#8666)

Browse Source 
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
Co-authored-by: Yineng Zhang <me@zhyncs.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
This commit is contained in:
blzheng
2026-05-21 13:09:18 +08:00
committed by GitHub
co-authored by JieXin Liang Yineng Zhang jianan-gu
parent 79b937aefb
commit 84ea47eb22
15 changed files with 481 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/decode.cpp
+83 -45
View File
@@ -1038,6 +1038,7 @@ void decode_attention_kernel_impl(
const index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
int64_t batches,
int64_t num_heads,
int64_t head_size,
@@ -1053,7 +1054,9 @@ void decode_attention_kernel_impl(
float logit_cap,
int64_t max_num_reqs,
int64_t max_context_len,
int64_t max_total_num_tokens) {
int64_t max_total_num_tokens,
bool is_cross_attn,
bool has_encoder_lens) {
using Vec = at::vec::Vectorized<float>;
// strides
@@ -1077,8 +1080,9 @@ void decode_attention_kernel_impl(
const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH;
// get key/value
int64_t seq_len_kv = seq_lens[bs];
int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
int64_t req_pool_id = req_pool_indices[bs];
int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
@@ -1102,7 +1106,7 @@ void decode_attention_kernel_impl(
/* A */ q_ptr,
/* B */ k_buffer + head_id * k_strideH,
/* C */ s_i,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* scl */ sm_scale,
/* M */ 1,
/* N */ n_size,
@@ -1142,7 +1146,7 @@ void decode_attention_kernel_impl(
/* A */ s_delta,
/* B */ v_buffer + head_id * v_strideH,
/* C */ v_prime,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* scl */ &m_delta,
/* M */ 1,
/* N */ head_size_v,
@@ -1159,6 +1163,8 @@ void decode_attention_kernel_impl(
at::vec::map<float>([s](Vec out) { return out * Vec(s); }, v_prime, v_prime, head_size_v);
v_prime[head_size_v] = m_prime + std::log(s_prime);
} else {
v_prime[head_size_v] = -std::numeric_limits<float>::infinity();
}
// move to the next index
@@ -1350,6 +1356,10 @@ void decode_attention_mla_kernel_impl(
[s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v);
(v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]);
}
} else {
for (int64_t h = 0; h < h_size; ++h) {
(v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits<float>::infinity();
}
}
// move to the next index
@@ -1372,6 +1382,7 @@ void decode_attention_grouped_kernel_impl(
const index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
int64_t batches,
int64_t num_heads,
int64_t num_heads_kv,
@@ -1388,7 +1399,9 @@ void decode_attention_grouped_kernel_impl(
float logit_cap,
int64_t max_num_reqs,
int64_t max_context_len,
int64_t max_total_num_tokens) {
int64_t max_total_num_tokens,
bool is_cross_attn,
bool has_encoder_lens) {
using Vec = at::vec::Vectorized<float>;
// block length for heads
@@ -1429,8 +1442,9 @@ void decode_attention_grouped_kernel_impl(
// get query
const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH;
int64_t seq_len_kv = seq_lens[bs];
int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
int64_t req_pool_id = req_pool_indices[bs];
int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
@@ -1456,7 +1470,7 @@ void decode_attention_grouped_kernel_impl(
/* A */ q_ptr,
/* B */ k_buffer + head_kv_id * k_strideH,
/* C */ s_i,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* scl */ sm_scale,
/* M */ h_size,
/* N */ n_size,
@@ -1500,7 +1514,7 @@ void decode_attention_grouped_kernel_impl(
/* A */ s_delta,
/* B */ v_buffer + head_kv_id * v_strideH,
/* C */ v_prime,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* scl */ m_delta,
/* M */ h_size,
/* N */ head_size_v,
@@ -1519,6 +1533,10 @@ void decode_attention_grouped_kernel_impl(
[s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v);
(v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]);
}
} else {
for (int64_t h = 0; h < h_size; ++h) {
(v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits<float>::infinity();
}
}
// move to the next index
@@ -1540,32 +1558,30 @@ void decode_attention_grouped_kernel_impl(
// req_to_token: [max_num_reqs, max_context_len] int32 or int64
// req_pool_indices: [num_seqs] int64
// seq_lens: [num_seqs] int64
// encoder_lens: [num_seqs] int64 or None
//
void decode_attention_cpu(
at::Tensor& query,
at::Tensor& k_buffer,
at::Tensor& v_buffer,
at::Tensor& output,
at::Tensor& key,
at::Tensor& value,
const std::optional<at::Tensor>& key,
const std::optional<at::Tensor>& value,
at::Tensor& loc,
at::Tensor& attn_logits,
at::Tensor& req_to_token,
at::Tensor& req_pool_indices,
at::Tensor& seq_lens,
double sm_scale,
double logit_cap) {
double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
// for MLA, key and value shares the same storage and value could be non-contiguous
CHECK_LAST_DIM_CONTIGUOUS_INPUT(key);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(value);
CHECK_DIM(3, query);
CHECK_DIM(3, k_buffer);
CHECK_DIM(3, v_buffer);
CHECK_DIM(3, key);
CHECK_DIM(3, value);
CHECK_DIM(1, loc);
int64_t num_seqs = seq_lens.size(0);
@@ -1580,7 +1596,6 @@ void decode_attention_cpu(
int64_t num_kv_splits = attn_logits.size(2);
CHECK_EQ(loc.numel(), num_seqs);
CHECK_EQ(attn_logits.size(0), num_seqs);
CHECK_EQ(attn_logits.size(1), num_heads);
CHECK_EQ(attn_logits.size(3), head_size_v + 1);
@@ -1595,11 +1610,6 @@ void decode_attention_cpu(
int64_t k_strideH = k_buffer.stride(1);
int64_t v_strideN = v_buffer.stride(0);
int64_t v_strideH = v_buffer.stride(1);
// strides for new key and value
int64_t nk_strideN = key.stride(0);
int64_t nk_strideH = key.stride(1);
int64_t nv_strideN = value.stride(0);
int64_t nv_strideH = value.stride(1);
// check index data types
const auto index_dtype = req_to_token.scalar_type();
@@ -1625,29 +1635,51 @@ void decode_attention_cpu(
int num_threads = at::get_num_threads();
int64_t size_per_thread = is_mla ? BLOCK_N * head_size + BLOCK_N * head_size_v : 0;
auto buffer = at::empty({num_threads, size_per_thread}, k_buffer.options());
bool has_encoder_lens = encoder_lens.has_value();
// Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
at::Tensor encoder_lens_t = seq_lens;
if (has_encoder_lens) {
encoder_lens_t = encoder_lens.value();
CHECK_EQ(encoder_lens_t.size(0), num_seqs);
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] {
// update the kv buffer
decode_set_kv_buffer(
(scalar_t*)k_buffer_data,
(scalar_t*)v_buffer_data,
key.data_ptr<scalar_t>(),
value.data_ptr<scalar_t>(),
loc.data_ptr<int64_t>(),
num_seqs,
num_heads_kv,
head_size,
head_size_v,
k_strideN,
k_strideH,
v_strideN,
v_strideH,
nk_strideN,
nk_strideH,
nv_strideN,
nv_strideH,
is_mla);
if (key.has_value()) {
TORCH_CHECK(value.has_value(), "key and value should have values at the same time")
CHECK_EQ(loc.numel(), num_seqs);
auto key_tensor = key.value();
auto value_tensor = value.value();
// for MLA, key and value shares the same storage and value could be non-contiguous
CHECK_LAST_DIM_CONTIGUOUS_INPUT(key_tensor);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(value_tensor);
CHECK_DIM(3, key_tensor);
CHECK_DIM(3, value_tensor);
// strides for new key and value
int64_t nk_strideN = key_tensor.stride(0);
int64_t nk_strideH = key_tensor.stride(1);
int64_t nv_strideN = value_tensor.stride(0);
int64_t nv_strideH = value_tensor.stride(1);
// update the kv buffer
decode_set_kv_buffer(
(scalar_t*)k_buffer_data,
(scalar_t*)v_buffer_data,
key_tensor.data_ptr<scalar_t>(),
value_tensor.data_ptr<scalar_t>(),
loc.data_ptr<int64_t>(),
num_seqs,
num_heads_kv,
head_size,
head_size_v,
k_strideN,
k_strideH,
v_strideN,
v_strideH,
nk_strideN,
nk_strideH,
nv_strideN,
nv_strideH,
is_mla);
}
if (num_heads == num_heads_kv) {
// MHA
@@ -1660,6 +1692,7 @@ void decode_attention_cpu(
req_to_token.data_ptr<index_t>(),
req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(),
encoder_lens_t.data_ptr<int64_t>(),
num_seqs,
num_heads,
head_size,
@@ -1675,7 +1708,9 @@ void decode_attention_cpu(
logit_cap,
max_num_reqs,
max_context_len,
max_total_num_tokens);
max_total_num_tokens,
is_cross_attn,
has_encoder_lens);
} else if (is_mla) {
// MLA
decode_attention_mla_kernel_impl<scalar_t, index_t, BLOCK_N>(
@@ -1716,6 +1751,7 @@ void decode_attention_cpu(
req_to_token.data_ptr<index_t>(),
req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(),
encoder_lens_t.data_ptr<int64_t>(),
num_seqs,
num_heads,
num_heads_kv,
@@ -1732,7 +1768,9 @@ void decode_attention_cpu(
logit_cap,
max_num_reqs,
max_context_len,
max_total_num_tokens);
max_total_num_tokens,
is_cross_attn,
has_encoder_lens);
}
});
});
sgl-kernel/csrc/cpu/extend.cpp
+112 -88
View File
@@ -22,6 +22,7 @@ void extend_attention_kernel_impl(
const index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
const index_t* __restrict__ extend_seq_lens,
const index_t* __restrict__ extend_start_loc,
const void* __restrict__ buffer,
@@ -46,7 +47,9 @@ void extend_attention_kernel_impl(
int max_total_num_tokens,
int max_len_extend,
int buffer_size_per_thread,
bool is_prefix_skipped) {
bool is_prefix_skipped,
bool is_cross_attn,
bool has_encoder_lens) {
// strides
const int o_strideM = num_heads * head_size_v;
const int o_strideH = head_size_v;
@@ -91,6 +94,7 @@ void extend_attention_kernel_impl(
int seq_extend_start_loc = extend_start_loc[bs];
int req_pool_id = req_pool_indices[bs];
int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!");
TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!");
TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
@@ -115,10 +119,11 @@ void extend_attention_kernel_impl(
fill_stub(v_prime, 0.f, m_size * head_size_v);
fill_stub(s_prime, 0.f, m_size);
fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size);
// stage 1: compute scores with prefix
for (int n = 0; n < seq_len_prefix; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, seq_len_prefix - n);
int kv_start = 0;
int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix;
for (int n = kv_start; n < kv_end; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, kv_end - n);
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
@@ -127,7 +132,7 @@ void extend_attention_kernel_impl(
pack_vnni<scalar_t, index_t>(
/* dst */ Btmp,
/* src */ k_buffer + head_kv_id * k_strideH,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ k_strideN,
@@ -153,7 +158,7 @@ void extend_attention_kernel_impl(
pack_vnni2<scalar_t, index_t>(
/* dst */ Btmp,
/* src */ v_buffer + head_kv_id * v_strideH,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ v_strideN,
@@ -172,92 +177,88 @@ void extend_attention_kernel_impl(
/* B */ Btmp,
/* C */ v_prime);
} // loop with seq_len_prefix
if (!is_cross_attn) {
// stage 2: compute the triangle part
int num_keys = std::min(seq_len_extend, m + BLOCK_M);
for (int n = 0; n < num_keys; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, num_keys - n);
// stage 2: compute the triangle part
int num_keys = std::min(seq_len_extend, m + BLOCK_M);
for (int n = 0; n < num_keys; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, num_keys - n);
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
// get key and pack
pack_vnni<scalar_t>(
/* dst */ Btmp,
/* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ ke_strideN,
/* ld_dst */ BLOCK_N);
// get key and pack
pack_vnni<scalar_t>(
/* dst */ Btmp,
/* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ ke_strideN,
/* ld_dst */ BLOCK_N);
// calculate s_i <- Q @ K
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ head_size,
/* lda */ q_strideM,
/* ldb */ BLOCK_N,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
// calculate s_i <- Q @ K
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ head_size,
/* lda */ q_strideM,
/* ldb */ BLOCK_N,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
// apply causal mask
// [Note] condition to apply causal mask.
// Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
// n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct
// when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. With
// BLOCK_M=512, BLOCK_N=768:
// BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
// Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
// Query row=0 is at position 512, so keys 513..767 are future and must be
// masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
// producing wrong (non-causal) attention for rows 0..254 of this m-block.
if (n + n_size - 1 > m) {
for (int row = 0; row < m_size; ++row) {
int last_col = m + row - n;
// [Note] mask the entire row if last_col < 0.
// Clamp to -1: when n > m + row every key in this block is a future
// key, so the entire row should be masked. Without this clamp,
// last_col+1 <= 0 and fill_stub would write before row_ptr.
// Example:
// For max_len_extend > 4096 → selects BLOCK_M=512, BLOCK_N=768
// m + BLOCK_M = 512 + 512 = 1024 > BLOCK_N = 768, this means we can have a a second n-block at n=768.
// For m = 512, row = 0, n = 768, last_col = 512 + 0 - 768 = -256 → out of bounds write in fill_stub
last_col = std::max(last_col, -1);
// fill [last_col + 1, n_size) to -inf
float* row_ptr = s_i + row * BLOCK_N;
fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
// apply causal mask
// [Note] condition to apply causal mask.
// Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
// n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct
// when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. With
// BLOCK_M=512, BLOCK_N=768:
// BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
// Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
// Query row=0 is at position 512, so keys 513..767 are future and must be
// masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
// producing wrong (non-causal) attention for rows 0..254 of this m-block.
if (n + n_size - 1 > m) {
for (int row = 0; row < m_size; ++row) {
int last_col = m + row - n;
// [Note] mask the entire row if last_col < 0.
// Clamp to -1: when n > m + row every key in this block is a future
// key, so the entire row should be masked. Without this clamp,
// last_col+1 <= 0 and fill_stub would write before row_ptr.
last_col = std::max(last_col, -1);
// fill [last_col + 1, n_size) to -inf
float* row_ptr = s_i + row * BLOCK_N;
fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
}
}
}
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale);
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale);
// get value and pack
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ ve_strideN,
/* ld_dst */ head_size_v);
// calculate V' <- s_delta @ V + V'
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ head_size_v,
/* K */ padded_n_size, // n_size
/* lda */ BLOCK_N,
/* ldb */ head_size_v,
/* ldc */ head_size_v,
/* add_C */ true,
/* A */ s_delta,
/* B */ Btmp,
/* C */ v_prime);
} // loop with seq_len_extend
// get value and pack
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ ve_strideN,
/* ld_dst */ head_size_v);
// calculate V' <- s_delta @ V + V'
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ head_size_v,
/* K */ padded_n_size, // n_size
/* lda */ BLOCK_N,
/* ldb */ head_size_v,
/* ldc */ head_size_v,
/* add_C */ true,
/* A */ s_delta,
/* B */ Btmp,
/* C */ v_prime);
} // loop with seq_len_extend
}
scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH;
for (int row = 0; row < m_size; ++row) {
float s = 1 / s_prime[row];
@@ -299,6 +300,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
req_to_token.data_ptr<index_t>(), \
req_pool_indices.data_ptr<int64_t>(), \
seq_lens.data_ptr<int64_t>(), \
encoder_lens_t.data_ptr<int64_t>(), \
extend_seq_lens.data_ptr<index_t>(), \
extend_start_loc.data_ptr<index_t>(), \
buffer.data_ptr(), \
@@ -323,7 +325,9 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
max_total_num_tokens, \
max_len_extend, \
sz, \
is_prefix_skipped); \
is_prefix_skipped, \
is_cross_attn, \
has_encoder_lens); \
} while (0)
// q_extend, k_extend, v_extend, o_extend: contiguous tensors
@@ -340,11 +344,12 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
// seq_lens: [num_seqs] int64
// extend_seq_lens: [num_seqs]
// extend_start_loc: [num_seqs]
// encoder_lens: [num_seqs] int64
//
void extend_attention_cpu(
at::Tensor& q_extend,
at::Tensor& k_extend,
at::Tensor& v_extend,
const std::optional<at::Tensor>& k_extend_opt,
const std::optional<at::Tensor>& v_extend_opt,
at::Tensor& o_extend,
at::Tensor& k_buffer,
at::Tensor& v_buffer,
@@ -355,7 +360,19 @@ void extend_attention_cpu(
at::Tensor& extend_start_loc,
int64_t max_len_extend,
double sm_scale,
double logit_cap) {
double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens) {
if (!is_cross_attn) {
TORCH_CHECK(
k_extend_opt.has_value() && v_extend_opt.has_value(),
"k_extend and v_extend are required for non-cross attention");
}
// Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer
// here.
auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer;
auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer;
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend);
CHECK_INPUT(o_extend);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend);
@@ -419,6 +436,13 @@ void extend_attention_cpu(
int num_threads = at::get_num_threads();
auto buffer = at::empty({}, q_extend.options().dtype(at::kChar));
bool has_encoder_lens = encoder_lens.has_value();
// Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
at::Tensor encoder_lens_t = seq_lens;
if (has_encoder_lens) {
encoder_lens_t = encoder_lens.value();
CHECK_EQ(encoder_lens_t.size(0), num_seqs);
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
if (max_len_extend <= 256) {
sgl-kernel/csrc/cpu/gemm.cpp
+14 -11
View File
@@ -722,10 +722,10 @@ at::Tensor convert_scale_packed(at::Tensor& scale) {
return packed_scale;
}
// mat1 : [M, K]
// mat1 : [*, K]
// mat2 : [N, K] ([K, N] if use_fma_gemm)
// bias : [N]
// out : [M, N]
// out : [*, N]
//
at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
@@ -735,23 +735,25 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
use_fma_gemm = true;
}
int64_t M = mat1.size(0);
int64_t K = mat1.size(1);
int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2);
CHECK_DIM(2, mat1);
const int64_t ndim = mat1.ndimension();
auto input_sizes = mat1.sizes().vec();
int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
int64_t K = use_fma_gemm ? mat1.size(1) : mat2.size(1);
int64_t M = use_fma_gemm ? mat1.size(0) : mat1.numel() / K;
CHECK_DIM(2, mat2);
if (!use_fma_gemm) {
CHECK_EQ(mat1.size(1), K);
if (use_fma_gemm) {
CHECK_DIM(2, mat1);
} else {
CHECK_EQ(mat1.size(ndim - 1), K);
}
auto dispatch_type = mat1.scalar_type();
auto out = at::empty({M, N}, mat1.options());
// strides
int64_t out_strideM = out.stride(0);
int64_t mat1_strideM = mat1.stride(0);
int64_t mat1_strideM = mat1.stride(-2);
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
@@ -787,7 +789,8 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
}
});
return out;
input_sizes[ndim - 1] = N;
return out.view(input_sizes);
}
// mat1 : [M, K]
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+16 -10
View File
@@ -90,20 +90,22 @@ void decode_attention_cpu(
at::Tensor& k_cache,
at::Tensor& v_cache,
at::Tensor& output,
at::Tensor& key,
at::Tensor& value,
const std::optional<at::Tensor>& key,
const std::optional<at::Tensor>& value,
at::Tensor& loc,
at::Tensor& attn_logits,
at::Tensor& req_to_token,
at::Tensor& req_pool_indices,
at::Tensor& seq_lens,
double sm_scale,
double logit_cap);
double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens);
void extend_attention_cpu(
at::Tensor& q_extend,
at::Tensor& k_extend,
at::Tensor& v_extend,
const std::optional<at::Tensor>& k_extend,
const std::optional<at::Tensor>& v_extend,
at::Tensor& o_extend,
at::Tensor& k_buffer,
at::Tensor& v_buffer,
@@ -114,7 +116,9 @@ void extend_attention_cpu(
at::Tensor& extend_start_loc,
int64_t max_len_extend,
double sm_scale,
double logit_cap);
double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens);
// flash attention
at::Tensor flash_attn_varlen_func(
@@ -462,16 +466,18 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// decode
m.def(
"decode_attention_cpu(Tensor query, Tensor k_cache, Tensor v_cahce, Tensor(a!) output, Tensor key, Tensor value, "
"decode_attention_cpu(Tensor query, Tensor k_cache, Tensor v_cahce, Tensor(a!) output, Tensor? key, Tensor? "
"value, "
"Tensor loc, Tensor attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, float sm_scale, "
"float logit_cap) -> ()");
"float logit_cap, bool is_cross_attn, Tensor? encoder_lens) -> ()");
m.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);
// extend
m.def(
"extend_attention_cpu(Tensor q_extend, Tensor k_extend, Tensor v_extend, Tensor(a!) o_extend, Tensor k_buffer, "
"extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? v_extend, Tensor(a!) o_extend, Tensor k_buffer, "
"Tensor v_buffer, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor extend_seq_lens, Tensor "
"extend_start_loc, int max_len_extend, float sm_scale, float logit_cap) -> ()");
"extend_start_loc, int max_len_extend, float sm_scale, float logit_cap, bool is_cross_attn, Tensor? "
"encoder_lens) -> ()");
m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
// flash attn