minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 17 Packages Projects Releases Wiki Activity

[NPU] Enable automatic ascend_attn selection for vision attention and graph runners (#31948)

Browse Source 
Co-authored-by: litao.dream <litao.dream@bytedance.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
Tao Li
2026-08-02 15:09:51 +08:00
committed by GitHub
co-authored by litao.dream Xinyuan Tong
parent a0b7bcf592
commit 8d106c3d79
6 changed files with 237 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py
+5 -4
View File
@@ -27,7 +27,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
)
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.runtime_context import get_mm
class ViTNpuGraphRunner(ViTCudaGraphRunner):
@@ -70,17 +69,19 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner):
graph = torch_npu.npu.NPUGraph()
vit = self.vit
override_backend = get_mm().mm_attention_backend
backend = self._attn_backend
with torch_npu.npu.graph(graph, pool=ViTNpuGraphRunner._graph_memory_pool):
y = None
deepstack_outs: List[torch.Tensor] = []
deepstack_capture_idx = 0
for layer_num, blk in enumerate(vit.blocks):
if override_backend == "ascend_attn":
if backend == "ascend_attn":
cu_seq_lens = self.cu_seq_lens[graph_key]
else:
raise RuntimeError("Not supported ViT attention backend")
raise RuntimeError(
f"ViT NPU graph does not support attention backend: {backend}"
)
if layer_num == 0:
y = blk(
python/sglang/srt/layers/attention/vision.py
+23 -1
View File
@@ -784,6 +784,9 @@ class VisionAscendAttention(nn.Module):
if not _is_npu:
raise Exception("VisionAscendAttention is only available for ascend npu")
super().__init__()
# Ascend fused attention does not support SGLang's additive masks, so
# masked inputs must stay on SDPA.
self.sdpa_fallback = VisionSdpaAttention(**kwargs)
def forward(
self,
@@ -795,6 +798,7 @@ class VisionAscendAttention(nn.Module):
seq_len: int,
softmax_scale: Optional[float] = None,
forward_metadata: Optional[VisionAttentionMetadata] = None,
attention_mask: Optional[torch.Tensor] = None,
**kwargs,
) -> torch.Tensor:
r"""
@@ -803,6 +807,19 @@ class VisionAscendAttention(nn.Module):
Returns:
[b * s, h, head_size]
"""
if attention_mask is not None:
return self.sdpa_fallback(
q=q,
k=k,
v=v,
cu_seqlens=cu_seqlens,
bsz=bsz,
seq_len=seq_len,
attention_mask=attention_mask,
forward_metadata=forward_metadata,
**kwargs,
)
if forward_metadata is not None:
# TND fused attention expects cumulative seqlens (cu_seqlens[1:]),
# not per-sequence lengths in forward_metadata.seq_lens.
@@ -1050,6 +1067,8 @@ class VisionAttention(nn.Module):
print_info_once(f"Multimodal attention backend not set. Use {qkv_backend}.")
print_info_once(f"Using {qkv_backend} as multimodal attention backend.")
self.qkv_backend_name: str = qkv_backend
self.customized_position_embedding_applier = (
customized_position_embedding_applier
)
@@ -1151,7 +1170,8 @@ class VisionAttention(nn.Module):
- CUDA (Hopper SM90): "fa3"
- CUDA (Blackwell SM100): "fa4"
- CUDA (other): "triton_attn"
- Non-CUDA: "sdpa"
- Ascend NPU: "ascend_attn"
- Other platforms: device-specific optimized backend or "sdpa"
"""
override_backend = get_mm().mm_attention_backend
if override_backend is not None:
@@ -1166,6 +1186,8 @@ class VisionAttention(nn.Module):
backend = "fa4"
else:
backend = "triton_attn"
elif _is_npu:
backend = "ascend_attn"
elif _is_musa:
if get_device_capability() >= (3, 1):
backend = "fa3"
python/sglang/srt/multimodal/internvl_vit_cuda_graph_runner.py
+13 -9
View File
@@ -22,7 +22,6 @@ import torch
import torch.nn as nn
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.runtime_context import get_mm
class InternViTCudaGraphRunner:
@@ -51,6 +50,7 @@ class InternViTCudaGraphRunner:
first_layer = encoder.layers[0]
# InternAttention wraps VisionAttention as first_layer.attn.attn
self._attn: VisionAttention = first_layer.attn.attn # type: ignore
self._attn_backend: str | None = getattr(self._attn, "qkv_backend_name", None)
@property
def device(self) -> torch.device:
@@ -95,17 +95,19 @@ class InternViTCudaGraphRunner:
def _warmup_once(self, key: Hashable) -> None:
"""Run a tiny eager warmup on the preallocated buffers to trigger lazy init."""
override_backend = get_mm().mm_attention_backend
backend = self._attn_backend
cu = self.cu[key]
cu_kk = self.cu_kk[key]
max_len = int(cu_kk.max().item()) if cu_kk.numel() else 0
if override_backend == "triton_attn":
if backend == "triton_attn":
cu_ws = [cu, cu_kk, max_len]
elif override_backend == "fa3":
elif backend == "fa3":
cu_ws = [cu, max_len]
else:
raise RuntimeError("Not supported ViT attention backend for InternVL CG")
raise RuntimeError(
f"InternVL ViT CUDA graph does not support attention backend: {backend}"
)
x = self.inp[key]
y = x
@@ -115,18 +117,20 @@ class InternViTCudaGraphRunner:
def _capture_graph(self, key: Hashable) -> None:
g = torch.cuda.CUDAGraph()
override_backend = get_mm().mm_attention_backend
backend = self._attn_backend
cu = self.cu[key]
cu_kk = self.cu_kk[key]
max_len = int(cu_kk.max().item()) if cu_kk.numel() else 0
if override_backend == "triton_attn":
if backend == "triton_attn":
cu_ws = [cu, cu_kk, max_len]
elif override_backend == "fa3":
elif backend == "fa3":
cu_ws = [cu, max_len]
else:
raise RuntimeError("Not supported ViT attention backend for InternVL CG")
raise RuntimeError(
f"InternVL ViT CUDA graph does not support attention backend: {backend}"
)
torch.cuda.synchronize()
python/sglang/srt/multimodal/vit_cuda_graph_runner.py
+10 -7
View File
@@ -25,7 +25,6 @@ import torch.nn as nn
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.runtime_context import get_mm
class ViTCudaGraphRunner:
@@ -80,7 +79,9 @@ class ViTCudaGraphRunner:
)
self._attn: Optional[VisionAttention] = getattr(first_blk, "attn", None)
self._attn_backend = getattr(self._attn, "qkv_backend", None)
self._attn_backend: Optional[str] = getattr(
self._attn, "qkv_backend_name", None
)
@property
def device(self) -> torch.device:
@@ -151,13 +152,13 @@ class ViTCudaGraphRunner:
cu_full_kk = self.cu_full_len_kk[graph_key]
max_full_len = int(cu_full_kk.max().item())
override_backend = get_mm().mm_attention_backend
backend = self._attn_backend
if self._fullatt_block_indexes and 0 not in vit.fullatt_block_indexes:
warmup_cu_ws = [cu_window, cu_window_kk, max_window_len]
else:
warmup_cu_ws = [cu_full, cu_full_kk, max_full_len]
if override_backend == "fa3":
if backend == "fa3":
warmup_cu_ws = [warmup_cu_ws[0], warmup_cu_ws[2]]
warmup_kwargs = dict(
@@ -192,12 +193,14 @@ class ViTCudaGraphRunner:
cu_seqlens_kk_now = cu_full_kk
max_len = max_full_len
if override_backend == "triton_attn":
if backend == "triton_attn":
cu_seq_len_ws = [cu_seqlens_now, cu_seqlens_kk_now, max_len]
elif override_backend == "fa3":
elif backend == "fa3":
cu_seq_len_ws = [cu_seqlens_now, max_len]
else:
raise RuntimeError("Not supported ViT attention backend")
raise RuntimeError(
f"ViT CUDA graph does not support attention backend: {backend}"
)
if position_embeddings is not None:
if layer_num == 0: