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
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)) o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else: else:
o = torch.empty_like(q) o = torch.empty_like(q)
cache_loc = (
if save_kv_cache: forward_batch.out_cache_loc
forward_batch.token_to_kv_pool.set_kv_buffer( if not layer.is_cross_attention
layer, forward_batch.out_cache_loc, k, v 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 _, max_extend_len = self.forward_metadata
self.extend_attention_fwd( self.extend_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k, k,
@@ -122,6 +123,8 @@ class IntelAMXAttnBackend(AttentionBackend):
max_extend_len, max_extend_len,
layer.scaling, layer.scaling,
layer.logit_cap, layer.logit_cap,
layer.is_cross_attention,
forward_batch.encoder_lens,
) )
return o 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)) o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else: else:
o = torch.empty_like(q) 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( self.decode_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), 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), o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
k, k,
v, v,
forward_batch.out_cache_loc, cache_loc,
attn_logits, attn_logits,
forward_batch.req_to_token_pool.req_to_token, forward_batch.req_to_token_pool.req_to_token,
forward_batch.req_pool_indices, forward_batch.req_pool_indices,
forward_batch.seq_lens, forward_batch.seq_lens,
layer.scaling, layer.scaling,
layer.logit_cap, layer.logit_cap,
layer.is_cross_attention,
forward_batch.encoder_lens,
) )
return o return o
def support_triton(self): 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 __future__ import annotations
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Optional
import torch import torch
from torch.nn.functional import scaled_dot_product_attention from torch.nn.functional import scaled_dot_product_attention
@@ -35,9 +35,11 @@ class TorchNativeAttnBackend(AttentionBackend):
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor, extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor, extend_seq_lens: torch.Tensor,
encoder_lens: Optional[torch.Tensor] = None,
scaling=None, scaling=None,
enable_gqa=False, enable_gqa=False,
causal=False, causal=False,
is_cross_attn=False,
): ):
"""Run the extend forward by using torch native sdpa op. """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] v_cache: [max_total_num_tokens, num_heads, head_size]
req_to_token: [max_num_reqs, max_context_len] req_to_token: [max_num_reqs, max_context_len]
req_pool_indices: [num_seqs] req_pool_indices: [num_seqs]
encoder_lens: [num_seqs] or None
seq_lens: [num_seqs] seq_lens: [num_seqs]
extend_prefix_lens: [num_seqs] extend_prefix_lens: [num_seqs]
extend_seq_lens: [num_seqs] extend_seq_lens: [num_seqs]
scaling: float or None scaling: float or None
enable_gqa: bool enable_gqa: bool
causal: bool causal: bool
is_cross_attn: bool
Returns: Returns:
output: [num_tokens, num_heads, head_size] output: [num_tokens, num_heads, head_size]
@@ -75,8 +79,16 @@ class TorchNativeAttnBackend(AttentionBackend):
seq_len_kv = seq_lens[seq_idx] seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q 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 = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty( per_req_query_redudant = torch.empty(
(per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]), (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 # get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence. # index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx] 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_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_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_to_token: torch.Tensor,
req_pool_indices: torch.Tensor, req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
encoder_lens: Optional[torch.Tensor] = None,
scaling=None, scaling=None,
enable_gqa=False, enable_gqa=False,
causal=False, causal=False,
is_cross_attn=False,
): ):
"""Run the decode forward by using torch native sdpa op. """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_to_token: [max_num_reqs, max_context_len]
req_pool_indices: [num_seqs] req_pool_indices: [num_seqs]
seq_lens: [num_seqs] seq_lens: [num_seqs]
encoder_lens: [num_seqs] or None
scaling: float or None scaling: float or None
enable_gqa: bool enable_gqa: bool
causal: bool causal: bool
is_cross_attn: bool
Returns: Returns:
output: [num_tokens, num_heads, head_size] output: [num_tokens, num_heads, head_size]
@@ -156,14 +172,24 @@ class TorchNativeAttnBackend(AttentionBackend):
seq_len_q = 1 seq_len_q = 1
seq_len_kv = seq_lens[seq_idx] seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q 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, :] per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache # get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence. # index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx] 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_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_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: else:
cache_loc = forward_batch.out_cache_loc 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) 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 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.seq_lens,
forward_batch.extend_prefix_lens, forward_batch.extend_prefix_lens,
forward_batch.extend_seq_lens, forward_batch.extend_seq_lens,
forward_batch.encoder_lens,
scaling=layer.scaling, scaling=layer.scaling,
enable_gqa=use_gqa, enable_gqa=use_gqa,
causal=causal, causal=causal,
is_cross_attn=layer.is_cross_attention,
) )
return o 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)) o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else: else:
o = torch.empty_like(q) 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: if layer.is_cross_attention:
cache_loc = forward_batch.encoder_out_cache_loc cache_loc = forward_batch.encoder_out_cache_loc
@@ -260,7 +293,8 @@ class TorchNativeAttnBackend(AttentionBackend):
cache_loc = forward_batch.out_cache_loc cache_loc = forward_batch.out_cache_loc
if save_kv_cache: 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 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_to_token_pool.req_to_token,
forward_batch.req_pool_indices, forward_batch.req_pool_indices,
forward_batch.seq_lens, forward_batch.seq_lens,
forward_batch.encoder_lens,
scaling=layer.scaling, scaling=layer.scaling,
enable_gqa=use_gqa, enable_gqa=use_gqa,
causal=False, causal=False,
is_cross_attn=layer.is_cross_attention,
) )
return o 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 = self.attn(q, k, v, forward_batch)
output = output.view(-1, self.num_local_heads * self.head_dim)
out, _ = self.o_proj(output) out, _ = self.o_proj(output)
return out return out
@@ -865,9 +866,7 @@ class MllamaForConditionalGeneration(nn.Module):
self.image_size, self.image_size,
dtype=torch.float32, dtype=torch.float32,
) )
batched_ar_ids = torch.ones( batched_ar_ids = torch.ones(bs, max_num_images, dtype=torch.int64)
bs, max_num_images, dtype=torch.int64, device="cuda"
)
batched_ar_mask = torch.zeros( batched_ar_mask = torch.zeros(
bs, max_num_images, max_num_tiles, dtype=torch.int64 bs, max_num_images, max_num_tiles, dtype=torch.int64
) )
@@ -886,11 +885,13 @@ class MllamaForConditionalGeneration(nn.Module):
img = pixel_values[0, j] img = pixel_values[0, j]
num_tiles = img.shape[0] num_tiles = img.shape[0]
batched_images[i, j, :num_tiles] = img 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[ batched_ar_mask[i, j, :num_tiles] = mm_input.mm_items[
0 0
].aspect_ratio_mask[0, j] ].model_specific_data["aspect_ratio_mask"][0, j]
i += 1 i += 1
return batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need 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 index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices, const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens, const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
int64_t batches, int64_t batches,
int64_t num_heads, int64_t num_heads,
int64_t head_size, int64_t head_size,
@@ -1053,7 +1054,9 @@ void decode_attention_kernel_impl(
float logit_cap, float logit_cap,
int64_t max_num_reqs, int64_t max_num_reqs,
int64_t max_context_len, 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>; using Vec = at::vec::Vectorized<float>;
// strides // strides
@@ -1077,8 +1080,9 @@ void decode_attention_kernel_impl(
const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH; const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH;
// get key/value // 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 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(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!"); 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, /* A */ q_ptr,
/* B */ k_buffer + head_id * k_strideH, /* B */ k_buffer + head_id * k_strideH,
/* C */ s_i, /* 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, /* scl */ sm_scale,
/* M */ 1, /* M */ 1,
/* N */ n_size, /* N */ n_size,
@@ -1142,7 +1146,7 @@ void decode_attention_kernel_impl(
/* A */ s_delta, /* A */ s_delta,
/* B */ v_buffer + head_id * v_strideH, /* B */ v_buffer + head_id * v_strideH,
/* C */ v_prime, /* 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, /* scl */ &m_delta,
/* M */ 1, /* M */ 1,
/* N */ head_size_v, /* 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); 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); 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 // 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); [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]); (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 // move to the next index
@@ -1372,6 +1382,7 @@ void decode_attention_grouped_kernel_impl(
const index_t* __restrict__ req_to_token, const index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices, const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens, const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
int64_t batches, int64_t batches,
int64_t num_heads, int64_t num_heads,
int64_t num_heads_kv, int64_t num_heads_kv,
@@ -1388,7 +1399,9 @@ void decode_attention_grouped_kernel_impl(
float logit_cap, float logit_cap,
int64_t max_num_reqs, int64_t max_num_reqs,
int64_t max_context_len, 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>; using Vec = at::vec::Vectorized<float>;
// block length for heads // block length for heads
@@ -1429,8 +1442,9 @@ void decode_attention_grouped_kernel_impl(
// get query // get query
const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH; 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 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(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!"); 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, /* A */ q_ptr,
/* B */ k_buffer + head_kv_id * k_strideH, /* B */ k_buffer + head_kv_id * k_strideH,
/* C */ s_i, /* 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, /* scl */ sm_scale,
/* M */ h_size, /* M */ h_size,
/* N */ n_size, /* N */ n_size,
@@ -1500,7 +1514,7 @@ void decode_attention_grouped_kernel_impl(
/* A */ s_delta, /* A */ s_delta,
/* B */ v_buffer + head_kv_id * v_strideH, /* B */ v_buffer + head_kv_id * v_strideH,
/* C */ v_prime, /* 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, /* scl */ m_delta,
/* M */ h_size, /* M */ h_size,
/* N */ head_size_v, /* 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); [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]); (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 // 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_to_token: [max_num_reqs, max_context_len] int32 or int64
// req_pool_indices: [num_seqs] int64 // req_pool_indices: [num_seqs] int64
// seq_lens: [num_seqs] int64 // seq_lens: [num_seqs] int64
// encoder_lens: [num_seqs] int64 or None
// //
void decode_attention_cpu( void decode_attention_cpu(
at::Tensor& query, at::Tensor& query,
at::Tensor& k_buffer, at::Tensor& k_buffer,
at::Tensor& v_buffer, at::Tensor& v_buffer,
at::Tensor& output, at::Tensor& output,
at::Tensor& key, const std::optional<at::Tensor>& key,
at::Tensor& value, const std::optional<at::Tensor>& value,
at::Tensor& loc, at::Tensor& loc,
at::Tensor& attn_logits, at::Tensor& attn_logits,
at::Tensor& req_to_token, at::Tensor& req_to_token,
at::Tensor& req_pool_indices, at::Tensor& req_pool_indices,
at::Tensor& seq_lens, at::Tensor& seq_lens,
double sm_scale, 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(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer); CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_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, query);
CHECK_DIM(3, k_buffer); CHECK_DIM(3, k_buffer);
CHECK_DIM(3, v_buffer); CHECK_DIM(3, v_buffer);
CHECK_DIM(3, key);
CHECK_DIM(3, value);
CHECK_DIM(1, loc); CHECK_DIM(1, loc);
int64_t num_seqs = seq_lens.size(0); 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); 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(0), num_seqs);
CHECK_EQ(attn_logits.size(1), num_heads); CHECK_EQ(attn_logits.size(1), num_heads);
CHECK_EQ(attn_logits.size(3), head_size_v + 1); 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 k_strideH = k_buffer.stride(1);
int64_t v_strideN = v_buffer.stride(0); int64_t v_strideN = v_buffer.stride(0);
int64_t v_strideH = v_buffer.stride(1); 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 // check index data types
const auto index_dtype = req_to_token.scalar_type(); const auto index_dtype = req_to_token.scalar_type();
@@ -1625,29 +1635,51 @@ void decode_attention_cpu(
int num_threads = at::get_num_threads(); int num_threads = at::get_num_threads();
int64_t size_per_thread = is_mla ? BLOCK_N * head_size + BLOCK_N * head_size_v : 0; 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()); 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_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] { AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] {
// update the kv buffer if (key.has_value()) {
decode_set_kv_buffer( TORCH_CHECK(value.has_value(), "key and value should have values at the same time")
(scalar_t*)k_buffer_data, CHECK_EQ(loc.numel(), num_seqs);
(scalar_t*)v_buffer_data, auto key_tensor = key.value();
key.data_ptr<scalar_t>(), auto value_tensor = value.value();
value.data_ptr<scalar_t>(), // for MLA, key and value shares the same storage and value could be non-contiguous
loc.data_ptr<int64_t>(), CHECK_LAST_DIM_CONTIGUOUS_INPUT(key_tensor);
num_seqs, CHECK_LAST_DIM_CONTIGUOUS_INPUT(value_tensor);
num_heads_kv, CHECK_DIM(3, key_tensor);
head_size, CHECK_DIM(3, value_tensor);
head_size_v, // strides for new key and value
k_strideN, int64_t nk_strideN = key_tensor.stride(0);
k_strideH, int64_t nk_strideH = key_tensor.stride(1);
v_strideN, int64_t nv_strideN = value_tensor.stride(0);
v_strideH, int64_t nv_strideH = value_tensor.stride(1);
nk_strideN, // update the kv buffer
nk_strideH, decode_set_kv_buffer(
nv_strideN, (scalar_t*)k_buffer_data,
nv_strideH, (scalar_t*)v_buffer_data,
is_mla); 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) { if (num_heads == num_heads_kv) {
// MHA // MHA
@@ -1660,6 +1692,7 @@ void decode_attention_cpu(
req_to_token.data_ptr<index_t>(), req_to_token.data_ptr<index_t>(),
req_pool_indices.data_ptr<int64_t>(), req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(), seq_lens.data_ptr<int64_t>(),
encoder_lens_t.data_ptr<int64_t>(),
num_seqs, num_seqs,
num_heads, num_heads,
head_size, head_size,
@@ -1675,7 +1708,9 @@ void decode_attention_cpu(
logit_cap, logit_cap,
max_num_reqs, max_num_reqs,
max_context_len, max_context_len,
max_total_num_tokens); max_total_num_tokens,
is_cross_attn,
has_encoder_lens);
} else if (is_mla) { } else if (is_mla) {
// MLA // MLA
decode_attention_mla_kernel_impl<scalar_t, index_t, BLOCK_N>( 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_to_token.data_ptr<index_t>(),
req_pool_indices.data_ptr<int64_t>(), req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(), seq_lens.data_ptr<int64_t>(),
encoder_lens_t.data_ptr<int64_t>(),
num_seqs, num_seqs,
num_heads, num_heads,
num_heads_kv, num_heads_kv,
@@ -1732,7 +1768,9 @@ void decode_attention_cpu(
logit_cap, logit_cap,
max_num_reqs, max_num_reqs,
max_context_len, 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 index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices, const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens, const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
const index_t* __restrict__ extend_seq_lens, const index_t* __restrict__ extend_seq_lens,
const index_t* __restrict__ extend_start_loc, const index_t* __restrict__ extend_start_loc,
const void* __restrict__ buffer, const void* __restrict__ buffer,
@@ -46,7 +47,9 @@ void extend_attention_kernel_impl(
int max_total_num_tokens, int max_total_num_tokens,
int max_len_extend, int max_len_extend,
int buffer_size_per_thread, int buffer_size_per_thread,
bool is_prefix_skipped) { bool is_prefix_skipped,
bool is_cross_attn,
bool has_encoder_lens) {
// strides // strides
const int o_strideM = num_heads * head_size_v; const int o_strideM = num_heads * head_size_v;
const int o_strideH = 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 seq_extend_start_loc = extend_start_loc[bs];
int req_pool_id = req_pool_indices[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_prefix >= 0, "prefix len < 0!");
TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!"); 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!"); 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(v_prime, 0.f, m_size * head_size_v);
fill_stub(s_prime, 0.f, m_size); fill_stub(s_prime, 0.f, m_size);
fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size); fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size);
// stage 1: compute scores with prefix // stage 1: compute scores with prefix
for (int n = 0; n < seq_len_prefix; n += BLOCK_N) { int kv_start = 0;
int n_size = std::min(BLOCK_N, seq_len_prefix - n); 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; // `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * 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>( pack_vnni<scalar_t, index_t>(
/* dst */ Btmp, /* dst */ Btmp,
/* src */ k_buffer + head_kv_id * k_strideH, /* 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, /* N */ n_size,
/* K */ head_size, /* K */ head_size,
/* ld_src */ k_strideN, /* ld_src */ k_strideN,
@@ -153,7 +158,7 @@ void extend_attention_kernel_impl(
pack_vnni2<scalar_t, index_t>( pack_vnni2<scalar_t, index_t>(
/* dst */ Btmp, /* dst */ Btmp,
/* src */ v_buffer + head_kv_id * v_strideH, /* 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, /* K */ n_size,
/* N */ head_size_v, /* N */ head_size_v,
/* ld_src */ v_strideN, /* ld_src */ v_strideN,
@@ -172,92 +177,88 @@ void extend_attention_kernel_impl(
/* B */ Btmp, /* B */ Btmp,
/* C */ v_prime); /* C */ v_prime);
} // loop with seq_len_prefix } // 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 // `n_size` is K in 2nd gemm, pad to TILE_K;
int num_keys = std::min(seq_len_extend, m + BLOCK_M); const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
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; // get key and pack
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; 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 // calculate s_i <- Q @ K
pack_vnni<scalar_t>( at::native::cpublas::brgemm(
/* dst */ Btmp, /* M */ m_size,
/* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH, /* N */ n_size,
/* N */ n_size, /* K */ head_size,
/* K */ head_size, /* lda */ q_strideM,
/* ld_src */ ke_strideN, /* ldb */ BLOCK_N,
/* ld_dst */ BLOCK_N); /* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
// calculate s_i <- Q @ K // apply causal mask
at::native::cpublas::brgemm( // [Note] condition to apply causal mask.
/* M */ m_size, // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
/* N */ n_size, // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct
/* K */ head_size, // when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. With
/* lda */ q_strideM, // BLOCK_M=512, BLOCK_N=768:
/* ldb */ BLOCK_N, // BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
/* ldc */ BLOCK_N, // Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
/* add_C */ false, // Query row=0 is at position 512, so keys 513..767 are future and must be
/* A */ q_ptr, // masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
/* B */ Btmp, // producing wrong (non-causal) attention for rows 0..254 of this m-block.
/* C */ s_i); if (n + n_size - 1 > m) {
for (int row = 0; row < m_size; ++row) {
// apply causal mask int last_col = m + row - n;
// [Note] condition to apply causal mask. // [Note] mask the entire row if last_col < 0.
// Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n + // Clamp to -1: when n > m + row every key in this block is a future
// n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct // key, so the entire row should be masked. Without this clamp,
// when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. With // last_col+1 <= 0 and fill_stub would write before row_ptr.
// BLOCK_M=512, BLOCK_N=768: last_col = std::max(last_col, -1);
// BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys. // fill [last_col + 1, n_size) to -inf
// Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768). float* row_ptr = s_i + row * BLOCK_N;
// Query row=0 is at position 512, so keys 513..767 are future and must be fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
// 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);
} }
}
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply( 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); 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 // get value and pack
pack_vnni2<scalar_t>( pack_vnni2<scalar_t>(
/* dst */ Btmp, /* dst */ Btmp,
/* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH, /* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
/* K */ n_size, /* K */ n_size,
/* N */ head_size_v, /* N */ head_size_v,
/* ld_src */ ve_strideN, /* ld_src */ ve_strideN,
/* ld_dst */ head_size_v); /* 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
// 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; 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) { for (int row = 0; row < m_size; ++row) {
float s = 1 / s_prime[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_to_token.data_ptr<index_t>(), \
req_pool_indices.data_ptr<int64_t>(), \ req_pool_indices.data_ptr<int64_t>(), \
seq_lens.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_seq_lens.data_ptr<index_t>(), \
extend_start_loc.data_ptr<index_t>(), \ extend_start_loc.data_ptr<index_t>(), \
buffer.data_ptr(), \ 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_total_num_tokens, \
max_len_extend, \ max_len_extend, \
sz, \ sz, \
is_prefix_skipped); \ is_prefix_skipped, \
is_cross_attn, \
has_encoder_lens); \
} while (0) } while (0)
// q_extend, k_extend, v_extend, o_extend: contiguous tensors // 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 // seq_lens: [num_seqs] int64
// extend_seq_lens: [num_seqs] // extend_seq_lens: [num_seqs]
// extend_start_loc: [num_seqs] // extend_start_loc: [num_seqs]
// encoder_lens: [num_seqs] int64
// //
void extend_attention_cpu( void extend_attention_cpu(
at::Tensor& q_extend, at::Tensor& q_extend,
at::Tensor& k_extend, const std::optional<at::Tensor>& k_extend_opt,
at::Tensor& v_extend, const std::optional<at::Tensor>& v_extend_opt,
at::Tensor& o_extend, at::Tensor& o_extend,
at::Tensor& k_buffer, at::Tensor& k_buffer,
at::Tensor& v_buffer, at::Tensor& v_buffer,
@@ -355,7 +360,19 @@ void extend_attention_cpu(
at::Tensor& extend_start_loc, at::Tensor& extend_start_loc,
int64_t max_len_extend, int64_t max_len_extend,
double sm_scale, 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_LAST_DIM_CONTIGUOUS_INPUT(q_extend);
CHECK_INPUT(o_extend); CHECK_INPUT(o_extend);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend); CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend);
@@ -419,6 +436,13 @@ void extend_attention_cpu(
int num_threads = at::get_num_threads(); int num_threads = at::get_num_threads();
auto buffer = at::empty({}, q_extend.options().dtype(at::kChar)); 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_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] { AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
if (max_len_extend <= 256) { 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; return packed_scale;
} }
// mat1 : [M, K] // mat1 : [*, K]
// mat2 : [N, K] ([K, N] if use_fma_gemm) // mat2 : [N, K] ([K, N] if use_fma_gemm)
// bias : [N] // bias : [N]
// out : [M, N] // out : [*, N]
// //
at::Tensor at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) { 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; 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_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2); 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); CHECK_DIM(2, mat2);
if (!use_fma_gemm) { if (use_fma_gemm) {
CHECK_EQ(mat1.size(1), K); CHECK_DIM(2, mat1);
} else {
CHECK_EQ(mat1.size(ndim - 1), K);
} }
auto dispatch_type = mat1.scalar_type(); auto dispatch_type = mat1.scalar_type();
auto out = at::empty({M, N}, mat1.options()); auto out = at::empty({M, N}, mat1.options());
// strides // strides
int64_t out_strideM = out.stride(0); 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 bool has_bias = bias.has_value();
const float* bias_data = nullptr; 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] // 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& k_cache,
at::Tensor& v_cache, at::Tensor& v_cache,
at::Tensor& output, at::Tensor& output,
at::Tensor& key, const std::optional<at::Tensor>& key,
at::Tensor& value, const std::optional<at::Tensor>& value,
at::Tensor& loc, at::Tensor& loc,
at::Tensor& attn_logits, at::Tensor& attn_logits,
at::Tensor& req_to_token, at::Tensor& req_to_token,
at::Tensor& req_pool_indices, at::Tensor& req_pool_indices,
at::Tensor& seq_lens, at::Tensor& seq_lens,
double sm_scale, double sm_scale,
double logit_cap); double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens);
void extend_attention_cpu( void extend_attention_cpu(
at::Tensor& q_extend, at::Tensor& q_extend,
at::Tensor& k_extend, const std::optional<at::Tensor>& k_extend,
at::Tensor& v_extend, const std::optional<at::Tensor>& v_extend,
at::Tensor& o_extend, at::Tensor& o_extend,
at::Tensor& k_buffer, at::Tensor& k_buffer,
at::Tensor& v_buffer, at::Tensor& v_buffer,
@@ -114,7 +116,9 @@ void extend_attention_cpu(
at::Tensor& extend_start_loc, at::Tensor& extend_start_loc,
int64_t max_len_extend, int64_t max_len_extend,
double sm_scale, double sm_scale,
double logit_cap); double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens);
// flash attention // flash attention
at::Tensor flash_attn_varlen_func( at::Tensor flash_attn_varlen_func(
@@ -462,16 +466,18 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// decode // decode
m.def( 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, " "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); m.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);
// extend // extend
m.def( 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 " "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); m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
// flash attn // 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_to_token: torch.Tensor,
req_pool_indices: torch.Tensor, req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None, scaling=None,
enable_gqa=False, enable_gqa=False,
causal=False, causal=False,
is_cross_attn=False,
): ):
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size] # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2) query = query.movedim(0, query.dim() - 2)
@@ -33,14 +35,21 @@ class TestDecodeAttention(CustomTestCase):
seq_len_q = 1 seq_len_q = 1
seq_len_kv = seq_lens[seq_idx] seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q 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, :] per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache # get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence. # index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx] 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_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_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 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 # This represents the number of tokens already in the sequence
seq_len = 1024 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) sm_scale = 1.0 / (D**0.5)
logit_cap = 0.0 logit_cap = 0.0
num_kv_splits = 8 num_kv_splits = 8
@@ -91,11 +103,12 @@ class TestDecodeAttention(CustomTestCase):
req_to_token = ( req_to_token = (
torch.arange(total_tokens, device=device) torch.arange(total_tokens, device=device)
.reshape(B, seq_len) .reshape(B, seq_len + encoder_len)
.to(torch.int32) .to(torch.int32)
) )
b_req_idx = torch.arange(B, device=device).to(torch.int64) b_req_idx = torch.arange(B, device=device).to(torch.int64)
b_seq_len = torch.full((B,), seq_len, 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( attn_logits = torch.empty(
(B, H_Q, num_kv_splits, D_V + 1), (B, H_Q, num_kv_splits, D_V + 1),
@@ -114,8 +127,8 @@ class TestDecodeAttention(CustomTestCase):
k_buffer, k_buffer,
v_buffer, v_buffer,
o, o,
key, key if not is_cross_attn else None,
value, value if not is_cross_attn else None,
loc, loc,
attn_logits, attn_logits,
req_to_token, req_to_token,
@@ -123,6 +136,8 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len, b_seq_len,
sm_scale, sm_scale,
logit_cap, logit_cap,
is_cross_attn,
encoder_lens,
) )
self._run_sdpa_forward_decode( self._run_sdpa_forward_decode(
@@ -135,15 +150,16 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len, b_seq_len,
scaling=sm_scale, scaling=sm_scale,
enable_gqa=enable_gqa, enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
) )
cos_sim = torch.nn.functional.cosine_similarity( cos_sim = torch.nn.functional.cosine_similarity(
o.flatten(), o_grouped.flatten(), dim=0 o.flatten(), o_grouped.flatten(), dim=0
) )
self.assertGreater(cos_sim.item(), 0.99) self.assertGreater(cos_sim.item(), 0.99)
torch.testing.assert_close(o, o_grouped, atol=3e-2, rtol=1e-6) 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 = [ configs = [
(2, 16, 16, 64, 64), (2, 16, 16, 64, 64),
(2, 16, 1, 16, 16), (2, 16, 1, 16, 16),
@@ -161,7 +177,10 @@ class TestDecodeAttention(CustomTestCase):
for B, H_Q, H_KV, D, D_V in configs: for B, H_Q, H_KV, D, D_V in configs:
for dtype in [torch.bfloat16, torch.float16]: for dtype in [torch.bfloat16, torch.float16]:
self._test_grouped_decode_attention_once( 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): 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, seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor, extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor, extend_seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None, scaling=None,
enable_gqa=False, enable_gqa=False,
causal=False, causal=False,
is_cross_attn=False,
): ):
assert seq_lens.shape[0] == extend_prefix_lens.shape[0] assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -43,7 +45,14 @@ class TestExtendAttention(CustomTestCase):
seq_len_kv = seq_lens[seq_idx] seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q 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 = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty( 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 # get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence. # index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx] 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_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_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, D,
DV, DV,
mla=False, mla=False,
is_cross_attn=False,
*, *,
b_seq_len_prefix=None, b_seq_len_prefix=None,
b_seq_len_extend=None, b_seq_len_extend=None,
@@ -94,32 +104,36 @@ class TestExtendAttention(CustomTestCase):
if b_seq_len_prefix is None: if b_seq_len_prefix is None:
b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32) b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
if mla:
b_seq_len_prefix.zero_()
else: else:
b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32) 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: if b_seq_len_extend is None:
b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32) b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
else: else:
b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32) 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 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) b_req_idx = torch.arange(B, dtype=torch.int32)
req_to_tokens = torch.empty((B, max_len_in_batch), 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 = 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 = torch.zeros((B,), dtype=torch.int32)
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
for i in range(B): for i in range(B):
req_to_tokens[i, : b_seq_len[i]] = torch.arange( 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] 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() extend_token_num = torch.sum(b_seq_len_extend).item()
H_BUF = 1 if mla else H_KV 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) q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
for i in range(B): for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] extend_start_in_buffer = (
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] 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_start = b_start_loc_extend[i]
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i] extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
k_extend[extend_start:extend_end] = k_buffer[ k_extend[extend_start:extend_end] = k_buffer[
@@ -178,7 +194,9 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend, b_seq_len_extend,
scaling=sm_scale, scaling=sm_scale,
enable_gqa=enable_gqa, 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) o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -197,16 +215,29 @@ class TestExtendAttention(CustomTestCase):
max_len_extend, max_len_extend,
sm_scale, sm_scale,
logit_cap, logit_cap,
is_cross_attn,
encoder_lens,
) )
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2) torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
def test_extend_attention(self): def test_extend_attention(self):
for is_mla in [True, False]: for is_mla in [True, False]:
self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla) for is_cross_attn in [True, False]:
self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla) if is_mla and is_cross_attn:
self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla) continue
self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla) 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): def test_extend_attention_large_seq_causal_mask(self):
self._test_extend_attention_once( self._test_extend_attention_once(
@@ -220,6 +251,18 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend=[5000], 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__": if __name__ == "__main__":
unittest.main() unittest.main()
test/registered/cpu/test_gemm.py
+10 -1
View File
@@ -35,6 +35,7 @@ class TestGemm(CustomTestCase):
N = [16, 32 * 13] N = [16, 32 * 13]
K = [32 * 16] K = [32 * 16]
has_bias = [False, True] has_bias = [False, True]
dim = [2, 3, 4, 5]
M_int8 = [2, 128] M_int8 = [2, 128]
N_int8 = [32 * 12] N_int8 = [32 * 12]
@@ -52,10 +53,16 @@ class TestGemm(CustomTestCase):
N_gptq = [4096] N_gptq = [4096]
K_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) mat1 = torch.randn(M, K, dtype=torch.bfloat16)
mat2 = torch.randn(N, 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()) ref = torch.matmul(mat1.float(), mat2.float().t())
if has_bias: if has_bias:
@@ -83,12 +90,14 @@ class TestGemm(CustomTestCase):
self.N, self.N,
self.K, self.K,
self.has_bias, self.has_bias,
self.dim,
): ):
with self.subTest( with self.subTest(
M=params[0], M=params[0],
N=params[1], N=params[1],
K=params[2], K=params[2],
has_bias=params[3], has_bias=params[3],
dim=params[4],
): ):
self._bf16_gemm(*params) self._bf16_gemm(*params)
test/registered/cpu/test_mla.py
+2
View File
@@ -118,6 +118,8 @@ class TestMLA(CustomTestCase):
b_seq_len, b_seq_len,
sm_scale, sm_scale,
logit_cap, logit_cap,
False,
None,
) )
self._run_sdpa_forward_decode( 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_to_token: torch.Tensor,
req_pool_indices: torch.Tensor, req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None, scaling=None,
enable_gqa=False, enable_gqa=False,
causal=False, causal=False,
is_cross_attn=False,
): ):
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size] # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2) query = query.movedim(0, query.dim() - 2)
@@ -30,14 +32,21 @@ class TestDecodeAttention(CustomTestCase):
seq_len_q = 1 seq_len_q = 1
seq_len_kv = seq_lens[seq_idx] seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q 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, :] per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache # get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence. # index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx] 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_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_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 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 # This represents the number of tokens already in the sequence
seq_len = 1024 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) sm_scale = 1.0 / (D**0.5)
logit_cap = 0.0 logit_cap = 0.0
num_kv_splits = 8 num_kv_splits = 8
@@ -88,11 +100,12 @@ class TestDecodeAttention(CustomTestCase):
req_to_token = ( req_to_token = (
torch.arange(total_tokens, device=device) torch.arange(total_tokens, device=device)
.reshape(B, seq_len) .reshape(B, seq_len + encoder_len)
.to(torch.int32) .to(torch.int32)
) )
b_req_idx = torch.arange(B, device=device).to(torch.int64) b_req_idx = torch.arange(B, device=device).to(torch.int64)
b_seq_len = torch.full((B,), seq_len, 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( attn_logits = torch.empty(
(B, H_Q, num_kv_splits, D_V + 1), (B, H_Q, num_kv_splits, D_V + 1),
@@ -111,8 +124,8 @@ class TestDecodeAttention(CustomTestCase):
k_buffer, k_buffer,
v_buffer, v_buffer,
o, o,
key, key if not is_cross_attn else None,
value, value if not is_cross_attn else None,
loc, loc,
attn_logits, attn_logits,
req_to_token, req_to_token,
@@ -120,6 +133,8 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len, b_seq_len,
sm_scale, sm_scale,
logit_cap, logit_cap,
is_cross_attn,
encoder_lens,
) )
self._run_sdpa_forward_decode( self._run_sdpa_forward_decode(
@@ -132,8 +147,9 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len, b_seq_len,
scaling=sm_scale, scaling=sm_scale,
enable_gqa=enable_gqa, enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
) )
cos_sim = torch.nn.functional.cosine_similarity( cos_sim = torch.nn.functional.cosine_similarity(
o.flatten(), o_grouped.flatten(), dim=0 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 B, H_Q, H_KV, D, D_V in configs:
for dtype in [torch.bfloat16, torch.float16]: for dtype in [torch.bfloat16, torch.float16]:
self._test_grouped_decode_attention_once( 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): 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, seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor, extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor, extend_seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None, scaling=None,
enable_gqa=False, enable_gqa=False,
causal=False, causal=False,
is_cross_attn=False,
): ):
assert seq_lens.shape[0] == extend_prefix_lens.shape[0] assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -40,7 +42,14 @@ class TestExtendAttention(CustomTestCase):
seq_len_kv = seq_lens[seq_idx] seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q 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 = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty( 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 # get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence. # index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx] 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_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_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, D,
DV, DV,
mla=False, mla=False,
is_cross_attn=False,
*, *,
b_seq_len_prefix=None, b_seq_len_prefix=None,
b_seq_len_extend=None, b_seq_len_extend=None,
@@ -91,32 +101,36 @@ class TestExtendAttention(CustomTestCase):
if b_seq_len_prefix is None: if b_seq_len_prefix is None:
b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32) b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
if mla:
b_seq_len_prefix.zero_()
else: else:
b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32) 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: if b_seq_len_extend is None:
b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32) b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
else: else:
b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32) 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 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) b_req_idx = torch.arange(B, dtype=torch.int32)
req_to_tokens = torch.empty((B, max_len_in_batch), 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 = 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 = torch.zeros((B,), dtype=torch.int32)
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
for i in range(B): for i in range(B):
req_to_tokens[i, : b_seq_len[i]] = torch.arange( 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] 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() extend_token_num = torch.sum(b_seq_len_extend).item()
H_BUF = 1 if mla else H_KV 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) q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
for i in range(B): for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] extend_start_in_buffer = (
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] 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_start = b_start_loc_extend[i]
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i] extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
k_extend[extend_start:extend_end] = k_buffer[ k_extend[extend_start:extend_end] = k_buffer[
@@ -175,7 +191,9 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend, b_seq_len_extend,
scaling=sm_scale, scaling=sm_scale,
enable_gqa=enable_gqa, 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) o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -194,16 +212,29 @@ class TestExtendAttention(CustomTestCase):
max_len_extend, max_len_extend,
sm_scale, sm_scale,
logit_cap, logit_cap,
is_cross_attn,
encoder_lens,
) )
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2) torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
def test_extend_attention(self): def test_extend_attention(self):
for is_mla in [True, False]: for is_mla in [True, False]:
self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla) for is_cross_attn in [True, False]:
self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla) if is_mla and is_cross_attn:
self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla) continue
self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla) 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): def test_extend_attention_large_seq_causal_mask(self):
self._test_extend_attention_once( 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] N = [16, 32 * 13]
K = [32 * 16] K = [32 * 16]
has_bias = [False, True] has_bias = [False, True]
dim = [2, 3, 4, 5]
M_int8 = [2, 128] M_int8 = [2, 128]
N_int8 = [32 * 12] N_int8 = [32 * 12]
@@ -49,10 +50,16 @@ class TestGemm(CustomTestCase):
N_gptq = [4096] N_gptq = [4096]
K_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) mat1 = torch.randn(M, K, dtype=torch.bfloat16)
mat2 = torch.randn(N, 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()) ref = torch.matmul(mat1.float(), mat2.float().t())
if has_bias: if has_bias:
@@ -80,12 +87,14 @@ class TestGemm(CustomTestCase):
self.N, self.N,
self.K, self.K,
self.has_bias, self.has_bias,
self.dim,
): ):
with self.subTest( with self.subTest(
M=params[0], M=params[0],
N=params[1], N=params[1],
K=params[2], K=params[2],
has_bias=params[3], has_bias=params[3],
dim=params[4],
): ):
self._bf16_gemm(*params) self._bf16_gemm(*params)
test/srt/cpu/test_mla.py
+2
View File
@@ -115,6 +115,8 @@ class TestMLA(CustomTestCase):
b_seq_len, b_seq_len,
sm_scale, sm_scale,
logit_cap, logit_cap,
False,
None,
) )
self._run_sdpa_forward_decode( self._run_sdpa_forward_decode(