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/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) {