minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
python/sglang/srt/layers/attention/intel_amx_backend.py
+17 -9
View File
@@ -99,14 +99,15 @@ class IntelAMXAttnBackend(AttentionBackend):
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else:
o = torch.empty_like(q)
if save_kv_cache:
forward_batch.token_to_kv_pool.set_kv_buffer(
layer, forward_batch.out_cache_loc, k, v
)
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
else forward_batch.encoder_out_cache_loc
)
if save_kv_cache and k is not None and v is not None:
forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
_, max_extend_len = self.forward_metadata
self.extend_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k,
@@ -122,6 +123,8 @@ class IntelAMXAttnBackend(AttentionBackend):
max_extend_len,
layer.scaling,
layer.logit_cap,
layer.is_cross_attention,
forward_batch.encoder_lens,
)
return o
@@ -142,7 +145,11 @@ class IntelAMXAttnBackend(AttentionBackend):
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else:
o = torch.empty_like(q)
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
else forward_batch.encoder_out_cache_loc
)
self.decode_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
@@ -150,15 +157,16 @@ class IntelAMXAttnBackend(AttentionBackend):
o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
k,
v,
forward_batch.out_cache_loc,
cache_loc,
attn_logits,
forward_batch.req_to_token_pool.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
layer.scaling,
layer.logit_cap,
layer.is_cross_attention,
forward_batch.encoder_lens,
)
return o
def support_triton(self):
python/sglang/srt/layers/attention/torch_native_backend.py
+44 -8
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
import torch
from torch.nn.functional import scaled_dot_product_attention
@@ -35,9 +35,11 @@ class TorchNativeAttnBackend(AttentionBackend):
seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
encoder_lens: Optional[torch.Tensor] = None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
"""Run the extend forward by using torch native sdpa op.
@@ -48,12 +50,14 @@ class TorchNativeAttnBackend(AttentionBackend):
v_cache: [max_total_num_tokens, num_heads, head_size]
req_to_token: [max_num_reqs, max_context_len]
req_pool_indices: [num_seqs]
encoder_lens: [num_seqs] or None
seq_lens: [num_seqs]
extend_prefix_lens: [num_seqs]
extend_seq_lens: [num_seqs]
scaling: float or None
enable_gqa: bool
causal: bool
is_cross_attn: bool
Returns:
output: [num_tokens, num_heads, head_size]
@@ -75,8 +79,16 @@ class TorchNativeAttnBackend(AttentionBackend):
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
if is_cross_attn:
start_kv = 0
end_kv = encoder_lens[seq_idx]
else:
start_kv = encoder_lens[seq_idx]
end_kv = start_kv + seq_len_kv
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty(
(per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]),
@@ -89,7 +101,7 @@ class TorchNativeAttnBackend(AttentionBackend):
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -123,9 +135,11 @@ class TorchNativeAttnBackend(AttentionBackend):
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
encoder_lens: Optional[torch.Tensor] = None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
"""Run the decode forward by using torch native sdpa op.
@@ -137,9 +151,11 @@ class TorchNativeAttnBackend(AttentionBackend):
req_to_token: [max_num_reqs, max_context_len]
req_pool_indices: [num_seqs]
seq_lens: [num_seqs]
encoder_lens: [num_seqs] or None
scaling: float or None
enable_gqa: bool
causal: bool
is_cross_attn: bool
Returns:
output: [num_tokens, num_heads, head_size]
@@ -156,14 +172,24 @@ class TorchNativeAttnBackend(AttentionBackend):
seq_len_q = 1
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
if is_cross_attn:
start_kv = 0
end_kv = encoder_lens[seq_idx]
else:
start_kv = encoder_lens[seq_idx]
end_kv = start_kv + seq_len_kv
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -208,7 +234,7 @@ class TorchNativeAttnBackend(AttentionBackend):
else:
cache_loc = forward_batch.out_cache_loc
if save_kv_cache:
if save_kv_cache and k is not None and v is not None:
forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
@@ -230,9 +256,11 @@ class TorchNativeAttnBackend(AttentionBackend):
forward_batch.seq_lens,
forward_batch.extend_prefix_lens,
forward_batch.extend_seq_lens,
forward_batch.encoder_lens,
scaling=layer.scaling,
enable_gqa=use_gqa,
causal=causal,
is_cross_attn=layer.is_cross_attention,
)
return o
@@ -253,6 +281,11 @@ class TorchNativeAttnBackend(AttentionBackend):
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else:
o = torch.empty_like(q)
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
else forward_batch.encoder_out_cache_loc
)
if layer.is_cross_attention:
cache_loc = forward_batch.encoder_out_cache_loc
@@ -260,7 +293,8 @@ class TorchNativeAttnBackend(AttentionBackend):
cache_loc = forward_batch.out_cache_loc
if save_kv_cache:
forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
if k is not None and v is not None:
forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
@@ -275,9 +309,11 @@ class TorchNativeAttnBackend(AttentionBackend):
forward_batch.req_to_token_pool.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
forward_batch.encoder_lens,
scaling=layer.scaling,
enable_gqa=use_gqa,
causal=False,
is_cross_attn=layer.is_cross_attention,
)
return o
python/sglang/srt/models/mllama.py
+6 -5
View File
@@ -565,6 +565,7 @@ class MllamaTextCrossAttention(nn.Module):
)
output = self.attn(q, k, v, forward_batch)
output = output.view(-1, self.num_local_heads * self.head_dim)
out, _ = self.o_proj(output)
return out
@@ -865,9 +866,7 @@ class MllamaForConditionalGeneration(nn.Module):
self.image_size,
dtype=torch.float32,
)
batched_ar_ids = torch.ones(
bs, max_num_images, dtype=torch.int64, device="cuda"
)
batched_ar_ids = torch.ones(bs, max_num_images, dtype=torch.int64)
batched_ar_mask = torch.zeros(
bs, max_num_images, max_num_tiles, dtype=torch.int64
)
@@ -886,11 +885,13 @@ class MllamaForConditionalGeneration(nn.Module):
img = pixel_values[0, j]
num_tiles = img.shape[0]
batched_images[i, j, :num_tiles] = img
batched_ar_ids[i, j] = mm_input.mm_items[0].aspect_ratio_ids[0, j]
batched_ar_ids[i, j] = mm_input.mm_items[0].model_specific_data[
"aspect_ratio_ids"
][0, j]
batched_ar_mask[i, j, :num_tiles] = mm_input.mm_items[
0
].aspect_ratio_mask[0, j]
].model_specific_data["aspect_ratio_mask"][0, j]
i += 1
return batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need
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
test/registered/cpu/test_decode.py
+29 -10
View File
@@ -21,9 +21,11 @@ class TestDecodeAttention(CustomTestCase):
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2)
@@ -33,14 +35,21 @@ class TestDecodeAttention(CustomTestCase):
seq_len_q = 1
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
end_kv = (
encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
)
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -61,10 +70,13 @@ class TestDecodeAttention(CustomTestCase):
return output
def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, dtype, device):
def _test_grouped_decode_attention_once(
self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
):
# This represents the number of tokens already in the sequence
seq_len = 1024
total_tokens = B * seq_len
encoder_len = 10
total_tokens = B * (seq_len + encoder_len)
sm_scale = 1.0 / (D**0.5)
logit_cap = 0.0
num_kv_splits = 8
@@ -91,11 +103,12 @@ class TestDecodeAttention(CustomTestCase):
req_to_token = (
torch.arange(total_tokens, device=device)
.reshape(B, seq_len)
.reshape(B, seq_len + encoder_len)
.to(torch.int32)
)
b_req_idx = torch.arange(B, device=device).to(torch.int64)
b_seq_len = torch.full((B,), seq_len, device=device).to(torch.int64)
encoder_lens = torch.full((B,), encoder_len, device=device).to(torch.int64)
attn_logits = torch.empty(
(B, H_Q, num_kv_splits, D_V + 1),
@@ -114,8 +127,8 @@ class TestDecodeAttention(CustomTestCase):
k_buffer,
v_buffer,
o,
key,
value,
key if not is_cross_attn else None,
value if not is_cross_attn else None,
loc,
attn_logits,
req_to_token,
@@ -123,6 +136,8 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len,
sm_scale,
logit_cap,
is_cross_attn,
encoder_lens,
)
self._run_sdpa_forward_decode(
@@ -135,15 +150,16 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len,
scaling=sm_scale,
enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
)
cos_sim = torch.nn.functional.cosine_similarity(
o.flatten(), o_grouped.flatten(), dim=0
)
self.assertGreater(cos_sim.item(), 0.99)
torch.testing.assert_close(o, o_grouped, atol=3e-2, rtol=1e-6)
def _test_grouped_decode_attention(self, device="cpu"):
def _test_grouped_decode_attention(self, device="cuda"):
configs = [
(2, 16, 16, 64, 64),
(2, 16, 1, 16, 16),
@@ -161,7 +177,10 @@ class TestDecodeAttention(CustomTestCase):
for B, H_Q, H_KV, D, D_V in configs:
for dtype in [torch.bfloat16, torch.float16]:
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, dtype=dtype, device=device
B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
)
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
)
def test_grouped_decode_attention(self):
test/registered/cpu/test_extend.py
+60 -17
View File
@@ -24,9 +24,11 @@ class TestExtendAttention(CustomTestCase):
seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -43,7 +45,14 @@ class TestExtendAttention(CustomTestCase):
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
end_kv = (
encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
)
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty(
@@ -57,7 +66,7 @@ class TestExtendAttention(CustomTestCase):
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -86,6 +95,7 @@ class TestExtendAttention(CustomTestCase):
D,
DV,
mla=False,
is_cross_attn=False,
*,
b_seq_len_prefix=None,
b_seq_len_extend=None,
@@ -94,32 +104,36 @@ class TestExtendAttention(CustomTestCase):
if b_seq_len_prefix is None:
b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
if mla:
b_seq_len_prefix.zero_()
else:
b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
if mla:
b_seq_len_prefix.zero_()
encoder_lens.zero_()
if b_seq_len_extend is None:
b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
else:
b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32)
b_seq_len = b_seq_len_prefix + b_seq_len_extend
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
max_len_in_batch = (
torch.max(b_seq_len, 0)[0].item() + torch.max(encoder_lens, 0)[0].item()
)
b_req_idx = torch.arange(B, dtype=torch.int32)
req_to_tokens = torch.empty((B, max_len_in_batch), dtype=torch.int32)
b_start_loc = torch.zeros((B,), dtype=torch.int32)
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1] + encoder_lens[:-1], 0)
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32)
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
for i in range(B):
req_to_tokens[i, : b_seq_len[i]] = torch.arange(
b_start_loc[i], b_start_loc[i] + b_seq_len[i]
req_to_tokens[i, : b_seq_len[i] + encoder_lens[i]] = torch.arange(
b_start_loc[i], b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
)
total_token_num = torch.sum(b_seq_len).item()
total_token_num = torch.sum(b_seq_len).item() + torch.sum(encoder_lens).item()
extend_token_num = torch.sum(b_seq_len_extend).item()
H_BUF = 1 if mla else H_KV
@@ -131,8 +145,10 @@ class TestExtendAttention(CustomTestCase):
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
extend_start_in_buffer = (
b_start_loc[i] + b_seq_len_prefix[i] + encoder_lens[i]
)
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
extend_start = b_start_loc_extend[i]
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
k_extend[extend_start:extend_end] = k_buffer[
@@ -178,7 +194,9 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend,
scaling=sm_scale,
enable_gqa=enable_gqa,
causal=True,
causal=not is_cross_attn,
is_cross_attn=is_cross_attn,
encoder_lens=encoder_lens,
)
o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -197,16 +215,29 @@ class TestExtendAttention(CustomTestCase):
max_len_extend,
sm_scale,
logit_cap,
is_cross_attn,
encoder_lens,
)
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
def test_extend_attention(self):
for is_mla in [True, False]:
self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla)
self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla)
self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla)
self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla)
for is_cross_attn in [True, False]:
if is_mla and is_cross_attn:
continue
self._test_extend_attention_once(
1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
)
def test_extend_attention_large_seq_causal_mask(self):
self._test_extend_attention_once(
@@ -220,6 +251,18 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend=[5000],
)
def test_extend_attention_gqa_partial_extend_with_prefix(self):
self._test_extend_attention_once(
B=1,
N_CTX=256,
H_Q=16,
H_KV=4,
D=128,
DV=96,
b_seq_len_prefix=[97],
b_seq_len_extend=[37],
)
if __name__ == "__main__":
unittest.main()
test/registered/cpu/test_gemm.py
+10 -1
View File
@@ -35,6 +35,7 @@ class TestGemm(CustomTestCase):
N = [16, 32 * 13]
K = [32 * 16]
has_bias = [False, True]
dim = [2, 3, 4, 5]
M_int8 = [2, 128]
N_int8 = [32 * 12]
@@ -52,10 +53,16 @@ class TestGemm(CustomTestCase):
N_gptq = [4096]
K_gptq = [4096]
def _bf16_gemm(self, M, N, K, has_bias):
def _bf16_gemm(self, M, N, K, has_bias, dim):
mat1 = torch.randn(M, K, dtype=torch.bfloat16)
mat2 = torch.randn(N, K, dtype=torch.bfloat16)
if dim == 3:
mat1 = mat1.unsqueeze(0).repeat(2, 1, 1)
if dim == 4:
mat1 = mat1.unsqueeze(0).unsqueeze(0).repeat(2, 2, 1, 1)
if dim == 5:
mat1 = mat1.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(2, 2, 2, 1, 1)
ref = torch.matmul(mat1.float(), mat2.float().t())
if has_bias:
@@ -83,12 +90,14 @@ class TestGemm(CustomTestCase):
self.N,
self.K,
self.has_bias,
self.dim,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
dim=params[4],
):
self._bf16_gemm(*params)
test/registered/cpu/test_mla.py
+2
View File
@@ -118,6 +118,8 @@ class TestMLA(CustomTestCase):
b_seq_len,
sm_scale,
logit_cap,
False,
None,
)
self._run_sdpa_forward_decode(
test/srt/cpu/test_decode.py
+28 -9
View File
@@ -18,9 +18,11 @@ class TestDecodeAttention(CustomTestCase):
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2)
@@ -30,14 +32,21 @@ class TestDecodeAttention(CustomTestCase):
seq_len_q = 1
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
end_kv = (
encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
)
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -58,10 +67,13 @@ class TestDecodeAttention(CustomTestCase):
return output
def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, dtype, device):
def _test_grouped_decode_attention_once(
self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
):
# This represents the number of tokens already in the sequence
seq_len = 1024
total_tokens = B * seq_len
encoder_len = 10
total_tokens = B * (seq_len + encoder_len)
sm_scale = 1.0 / (D**0.5)
logit_cap = 0.0
num_kv_splits = 8
@@ -88,11 +100,12 @@ class TestDecodeAttention(CustomTestCase):
req_to_token = (
torch.arange(total_tokens, device=device)
.reshape(B, seq_len)
.reshape(B, seq_len + encoder_len)
.to(torch.int32)
)
b_req_idx = torch.arange(B, device=device).to(torch.int64)
b_seq_len = torch.full((B,), seq_len, device=device).to(torch.int64)
encoder_lens = torch.full((B,), encoder_len, device=device).to(torch.int64)
attn_logits = torch.empty(
(B, H_Q, num_kv_splits, D_V + 1),
@@ -111,8 +124,8 @@ class TestDecodeAttention(CustomTestCase):
k_buffer,
v_buffer,
o,
key,
value,
key if not is_cross_attn else None,
value if not is_cross_attn else None,
loc,
attn_logits,
req_to_token,
@@ -120,6 +133,8 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len,
sm_scale,
logit_cap,
is_cross_attn,
encoder_lens,
)
self._run_sdpa_forward_decode(
@@ -132,8 +147,9 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len,
scaling=sm_scale,
enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
)
cos_sim = torch.nn.functional.cosine_similarity(
o.flatten(), o_grouped.flatten(), dim=0
)
@@ -158,7 +174,10 @@ class TestDecodeAttention(CustomTestCase):
for B, H_Q, H_KV, D, D_V in configs:
for dtype in [torch.bfloat16, torch.float16]:
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, dtype=dtype, device=device
B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
)
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
)
def test_grouped_decode_attention(self):
test/srt/cpu/test_extend.py
+48 -17
View File
@@ -21,9 +21,11 @@ class TestExtendAttention(CustomTestCase):
seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -40,7 +42,14 @@ class TestExtendAttention(CustomTestCase):
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
end_kv = (
encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
)
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty(
@@ -54,7 +63,7 @@ class TestExtendAttention(CustomTestCase):
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -83,6 +92,7 @@ class TestExtendAttention(CustomTestCase):
D,
DV,
mla=False,
is_cross_attn=False,
*,
b_seq_len_prefix=None,
b_seq_len_extend=None,
@@ -91,32 +101,36 @@ class TestExtendAttention(CustomTestCase):
if b_seq_len_prefix is None:
b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
if mla:
b_seq_len_prefix.zero_()
else:
b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
if mla:
b_seq_len_prefix.zero_()
encoder_lens.zero_()
if b_seq_len_extend is None:
b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
else:
b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32)
b_seq_len = b_seq_len_prefix + b_seq_len_extend
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
max_len_in_batch = (
torch.max(b_seq_len, 0)[0].item() + torch.max(encoder_lens, 0)[0].item()
)
b_req_idx = torch.arange(B, dtype=torch.int32)
req_to_tokens = torch.empty((B, max_len_in_batch), dtype=torch.int32)
b_start_loc = torch.zeros((B,), dtype=torch.int32)
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1] + encoder_lens[:-1], 0)
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32)
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
for i in range(B):
req_to_tokens[i, : b_seq_len[i]] = torch.arange(
b_start_loc[i], b_start_loc[i] + b_seq_len[i]
req_to_tokens[i, : b_seq_len[i] + encoder_lens[i]] = torch.arange(
b_start_loc[i], b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
)
total_token_num = torch.sum(b_seq_len).item()
total_token_num = torch.sum(b_seq_len).item() + torch.sum(encoder_lens).item()
extend_token_num = torch.sum(b_seq_len_extend).item()
H_BUF = 1 if mla else H_KV
@@ -128,8 +142,10 @@ class TestExtendAttention(CustomTestCase):
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
extend_start_in_buffer = (
b_start_loc[i] + b_seq_len_prefix[i] + encoder_lens[i]
)
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
extend_start = b_start_loc_extend[i]
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
k_extend[extend_start:extend_end] = k_buffer[
@@ -175,7 +191,9 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend,
scaling=sm_scale,
enable_gqa=enable_gqa,
causal=True,
causal=not is_cross_attn,
is_cross_attn=is_cross_attn,
encoder_lens=encoder_lens,
)
o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -194,16 +212,29 @@ class TestExtendAttention(CustomTestCase):
max_len_extend,
sm_scale,
logit_cap,
is_cross_attn,
encoder_lens,
)
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
def test_extend_attention(self):
for is_mla in [True, False]:
self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla)
self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla)
self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla)
self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla)
for is_cross_attn in [True, False]:
if is_mla and is_cross_attn:
continue
self._test_extend_attention_once(
1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
)
def test_extend_attention_large_seq_causal_mask(self):
self._test_extend_attention_once(
test/srt/cpu/test_gemm.py
+10 -1
View File
@@ -32,6 +32,7 @@ class TestGemm(CustomTestCase):
N = [16, 32 * 13]
K = [32 * 16]
has_bias = [False, True]
dim = [2, 3, 4, 5]
M_int8 = [2, 128]
N_int8 = [32 * 12]
@@ -49,10 +50,16 @@ class TestGemm(CustomTestCase):
N_gptq = [4096]
K_gptq = [4096]
def _bf16_gemm(self, M, N, K, has_bias):
def _bf16_gemm(self, M, N, K, has_bias, dim):
mat1 = torch.randn(M, K, dtype=torch.bfloat16)
mat2 = torch.randn(N, K, dtype=torch.bfloat16)
if dim == 3:
mat1 = mat1.unsqueeze(0).repeat(2, 1, 1)
if dim == 4:
mat1 = mat1.unsqueeze(0).unsqueeze(0).repeat(2, 2, 1, 1)
if dim == 5:
mat1 = mat1.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(2, 2, 2, 1, 1)
ref = torch.matmul(mat1.float(), mat2.float().t())
if has_bias:
@@ -80,12 +87,14 @@ class TestGemm(CustomTestCase):
self.N,
self.K,
self.has_bias,
self.dim,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
dim=params[4],
):
self._bf16_gemm(*params)
test/srt/cpu/test_mla.py
+2
View File
@@ -115,6 +115,8 @@ class TestMLA(CustomTestCase):
b_seq_len,
sm_scale,
logit_cap,
False,
None,
)
self._run_sdpa_forward_decode(