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