From e86476acfc525d7c378efbeab2321d5bc65b7899 Mon Sep 17 00:00:00 2001 From: JiaruiChang5268 <71279366+JiaruiChang5268@users.noreply.github.com> Date: Sat, 31 Jan 2026 08:49:38 +0800 Subject: [PATCH] [NPU] support llama-3.2-11B-vision-instruct mode for NPU (#17492) Co-authored-by: McZyWu Co-authored-by: chenyang08056032 Co-authored-by: Hexq0210 <893781835@qq.com> --- .../npu/attention/ascend_backend.py | 75 ++++++- .../attention/ascend_torch_native_backend.py | 201 ++++++++++++++++++ python/sglang/srt/layers/linear.py | 4 +- ...st_ascend_llama_3_2_11b_vision_instruct.py | 29 +++ 4 files changed, 296 insertions(+), 13 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py create mode 100644 test/registered/ascend/vlm_models/test_ascend_llama_3_2_11b_vision_instruct.py diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 81e1241ee..f18c57070 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -11,13 +11,15 @@ from sgl_kernel_npu.attention.sinks_attention import ( ) from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.hardware_backend.npu.attention.ascend_torch_native_backend import ( + AscendTorchNativeAttnBackend, +) from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( is_fia_nz, is_mla_preprocess_enabled, ) from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp -from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.speculative.spec_info import SpecInput @@ -223,7 +225,7 @@ class AscendAttnBackend(AttentionBackend): self.q_head_dim = self.qk_rope_head_dim + self.qk_nope_head_dim else: self.use_alibi = getattr(model_runner.model_config, "use_alibi", False) - self.native_attn = TorchNativeAttnBackend(model_runner) + self.native_attn = AscendTorchNativeAttnBackend() self.graph_metadata = {} self.max_context_len = model_runner.model_config.context_len self.req_to_token = model_runner.req_to_token_pool.req_to_token @@ -751,10 +753,15 @@ class AscendAttnBackend(AttentionBackend): ) if not self.use_mla: - if save_kv_cache: - forward_batch.token_to_kv_pool.set_kv_buffer( - layer, forward_batch.out_cache_loc, k, v + # In cross attention layer, when there is no vision input,the values of k and v is None + if save_kv_cache and k is not None and v is not None: + # support cross attention + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc ) + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) @@ -812,7 +819,13 @@ class AscendAttnBackend(AttentionBackend): ): causal = False - if layer.qk_head_dim <= 128 and causal: + # there are some accuracy issues in cross attention scene to use torch_npu._npu_flash_attention_qlens + # forward_batch.encoder_lens is not None in cross attention scend, we add native attn to solve accuracy issues + if ( + layer.qk_head_dim <= 128 + and causal + and forward_batch.encoder_lens is None + ): if not self.use_alibi: query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) attn_output = torch.empty( @@ -860,7 +873,8 @@ class AscendAttnBackend(AttentionBackend): q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) o_ = attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim) - self.native_attn._run_sdpa_forward_extend( + # add forward_batch.encoder_lens and is_cross_attention arguments for cross attention scene + attn_output = self.native_attn.run_sdpa_forward_extend( q_, o_, k_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim), @@ -870,10 +884,15 @@ class AscendAttnBackend(AttentionBackend): forward_batch.seq_lens, forward_batch.extend_prefix_lens, forward_batch.extend_seq_lens, + forward_batch.encoder_lens, + is_cross_attention=layer.is_cross_attention, scaling=layer.scaling, enable_gqa=use_gqa, causal=causal, ) + attn_output = attn_output.view( + -1, layer.tp_q_head_num * layer.v_head_dim + ) elif sum(forward_batch.extend_prefix_lens_cpu) > 0: num_token_padding = q.shape[0] q, k, v = [ @@ -1440,10 +1459,15 @@ class AscendAttnBackend(AttentionBackend): ) if not self.use_mla: - if save_kv_cache: - forward_batch.token_to_kv_pool.set_kv_buffer( - layer, forward_batch.out_cache_loc, k, v + # In cross attention layer, when there is no vision input,the values of k and v is None + if save_kv_cache and k is not None and v is not None: + # support cross attention + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc ) + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) num_tokens = q.shape[0] k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) @@ -1492,7 +1516,9 @@ class AscendAttnBackend(AttentionBackend): actual_seq_lengths_kv=actual_seq_len_kv, scale=layer.scaling, ) - else: + # there are some accuracy issues in cross attention scene to use torch_npu._npu_flash_attention_qlens + # forward_batch.encoder_lens is not None in cross attention scend, we add native attn to solve accuracy issues + elif forward_batch.encoder_lens is None: query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) num_tokens = query.shape[0] if not self.use_alibi: @@ -1526,6 +1552,33 @@ class AscendAttnBackend(AttentionBackend): slopes=slopes, is_extend=False, ) + else: + if layer.qk_head_dim != layer.v_head_dim: + attn_output = q.new_empty( + (q.shape[0], layer.tp_q_head_num * layer.v_head_dim) + ) + else: + attn_output = torch.empty_like(q) + + use_gqa = layer.tp_q_head_num != layer.tp_k_head_num + + q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + o_ = attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + attn_output = self.native_attn.run_sdpa_forward_decode( + q_, + o_, + k_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim), + v_cache.view(-1, layer.tp_v_head_num, layer.v_head_dim), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.encoder_lens, + is_cross_attention=layer.is_cross_attention, + scaling=layer.scaling, + enable_gqa=use_gqa, + causal=False, + ) return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim) else: if save_kv_cache: diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py new file mode 100644 index 000000000..4c3b43d69 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +import torch +from torch.nn.functional import scaled_dot_product_attention + + +class AscendTorchNativeAttnBackend: + def __init__(self): + pass + + def run_sdpa_forward_extend( + self, + query: torch.Tensor, + output: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + encoder_lens: torch.Tensor = None, + is_cross_attention: bool = False, + scaling=None, + enable_gqa=False, + causal=False, + ): + """Run the extend forward by using torch native sdpa op. + + Args: + query: [num_tokens, num_heads, head_size] + output: [num_tokens, num_heads, head_size] + k_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_pool_indices: [num_seqs] + seq_lens: [num_seqs] + extend_prefix_lens: [num_seqs] + extend_seq_lens: [num_seqs] + encoder_lens: [num_seqs] + is_cross_attention: [bool] + scaling: float or None + enable_gqa: bool + causal: bool + + Returns: + output: [num_tokens, num_heads, head_size] + """ + + assert seq_lens.shape[0] == extend_prefix_lens.shape[0] + assert seq_lens.shape[0] == extend_seq_lens.shape[0] + + # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size] + query = query.movedim(0, query.dim() - 2) + + start_q, start_kv = 0, 0 + for seq_idx in range(seq_lens.shape[0]): + # Need optimize the performance later. + + extend_seq_len_q = extend_seq_lens[seq_idx] + prefill_seq_len_q = extend_prefix_lens[seq_idx] + + seq_len_kv = seq_lens[seq_idx] + end_q = start_q + extend_seq_len_q + end_kv = start_kv + seq_len_kv + atten_start_kv = 0 + atten_end_kv = seq_lens[seq_idx] + # support cross attention + if encoder_lens is not None: + if is_cross_attention: + atten_end_kv = encoder_lens[seq_idx] + else: + atten_start_kv = encoder_lens[seq_idx] + atten_end_kv = encoder_lens[seq_idx] + extend_seq_len_q + + 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]), + dtype=per_req_query.dtype, + device=per_req_query.device, + ) + + per_req_query_redudant[:, prefill_seq_len_q:, :] = per_req_query + + # 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, atten_start_kv:atten_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) + + if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype): + # scaled_dot_product_attention() expects query, key, and value to have the same dtype + per_req_key = per_req_key.to(per_req_query.dtype) + per_req_value = per_req_value.to(per_req_query.dtype) + + per_req_out_redudant = ( + scaled_dot_product_attention( + per_req_query_redudant.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) + ) + output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :] + start_q, start_kv = end_q, end_kv + return output + + def run_sdpa_forward_decode( + self, + query: torch.Tensor, + output: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: torch.Tensor = None, + is_cross_attention: bool = False, + scaling=None, + enable_gqa=False, + causal=False, + ): + """Run the decode forward by using torch native sdpa op. + + Args: + query: [num_tokens, num_heads, head_size] + output: [num_tokens, num_heads, head_size] + k_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_pool_indices: [num_seqs] + seq_lens: [num_seqs] + encoder_lens: [num_seqs] + is_cross_attention: [bool] + scaling: float or None + enable_gqa: bool + causal: bool + + Returns: + output: [num_tokens, num_heads, head_size] + """ + + # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size] + query = query.movedim(0, query.dim() - 2) + + start_q, start_kv = 0, 0 + for seq_idx in range(seq_lens.shape[0]): + # Need optimize the performance later. + + 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 + atten_start_kv = 0 + atten_end_kv = seq_lens[seq_idx] + # support cross attention + if encoder_lens is not None: + if is_cross_attention: + atten_end_kv = encoder_lens[seq_idx] + else: + atten_start_kv = encoder_lens[seq_idx] + atten_end_kv = encoder_lens[seq_idx] + 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, atten_start_kv:atten_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) + + if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype): + # scaled_dot_product_attention() expects query, key, and value to have the same dtype + per_req_key = per_req_key.to(per_req_query.dtype) + per_req_value = per_req_value.to(per_req_query.dtype) + + per_req_out = ( + scaled_dot_product_attention( + per_req_query.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) + ) + output[start_q:end_q, :, :] = per_req_out + start_q, start_kv = end_q, end_kv + + return output + + def support_triton(self): + return False diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index ebe7bb581..90ff6dc71 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -354,7 +354,7 @@ class ColumnParallelLinear(LinearBase): ) if bias: self.bias = Parameter( - torch.empty(self.output_size_per_partition, dtype=params_dtype) + torch.zeros(self.output_size_per_partition, dtype=params_dtype) ) set_weight_attrs( self.bias, @@ -1302,7 +1302,7 @@ class RowParallelLinear(LinearBase): ) if bias: - self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype)) + self.bias = Parameter(torch.zeros(self.output_size, dtype=params_dtype)) set_weight_attrs( self.bias, { diff --git a/test/registered/ascend/vlm_models/test_ascend_llama_3_2_11b_vision_instruct.py b/test/registered/ascend/vlm_models/test_ascend_llama_3_2_11b_vision_instruct.py new file mode 100644 index 000000000..eba8dec56 --- /dev/null +++ b/test/registered/ascend/vlm_models/test_ascend_llama_3_2_11b_vision_instruct.py @@ -0,0 +1,29 @@ +import unittest + +from sglang.test.ascend.vlm_utils import TestVLMModels +from sglang.test.ci.ci_register import register_npu_ci + +register_npu_ci(est_time=400, suite="nightly-1-npu-a3", nightly=True) + + +class TestLlama3211BVisionInstruct(TestVLMModels): + model = ( + "/root/.cache/modelscope/hub/models/LLM-Research/Llama-3.2-11B-Vision-Instruct" + ) + mmmu_accuracy = 0.2 + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + "0.8", + "--attention-backend", + "ascend", + "--disable-cuda-graph", + "--disable-radix-cache", + ] + + def test_vlm_mmmu_benchmark(self): + self._run_vlm_mmmu_test() + + +if __name__ == "__main__": + unittest.main()