chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Convert cu_seqlens to CPU for npu_flash_attention_unpad operator (#15434)

Browse Source
This commit is contained in:
fy
2026-01-04 08:16:42 +08:00
committed by GitHub
parent ef5ac6f01e
commit 25fa2ac290
9 changed files with 36 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/attention/vision.py
View File

@@ -470,7 +470,7 @@ class VisionAscendAttention(nn.Module):
Returns:
[b * s, h, head_size]
"""
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device="cpu")
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
if seq_lens.is_npu:

5
python/sglang/srt/models/dots_vlm_vit.py
View File

@@ -12,7 +12,7 @@ from sglang.srt.configs.dots_vlm import DotsVisionConfig
from sglang.srt.distributed import parallel_state
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.quantization import QuantizationConfig
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__)
@@ -315,6 +315,9 @@ class DotsVisionTransformer(PreTrainedModel):
dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
)
cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu():
cu_seqlens = cu_seqlens.to("cpu")
for blk in self.blocks:
hidden_states = blk(

6
python/sglang/srt/models/glm4v.py
View File

@@ -57,7 +57,7 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.glm4 import Glm4Model
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
from sglang.srt.utils.hf_transformers_utils import get_processor
logger = logging.getLogger(__name__)
@@ -514,6 +514,10 @@ class Glm4vVisionModel(nn.Module):
rotary_pos_emb_cos = torch.cat([rotary_pos_emb_cos, rotary_pos_emb_cos], dim=-1)
rotary_pos_emb_sin = torch.cat([rotary_pos_emb_sin, rotary_pos_emb_sin], dim=-1)
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu():
cu_seqlens = cu_seqlens.to("cpu")
# x.shape: (s, b, d) where b=1 for vision processing
# transformers
x = x.unsqueeze(1)

5
python/sglang/srt/models/idefics2.py
View File

@@ -28,7 +28,7 @@ from sglang.srt.layers.activation import get_act_fn
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
class Idefics2VisionMLP(nn.Module):
@@ -161,6 +161,9 @@ class Idefics2Encoder(nn.Module):
`input_ids` indices into associated vectorsthan the model's
internal embedding lookup matrix.
"""
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu():
cu_seqlens = cu_seqlens.to("cpu")
hidden_states = inputs_embeds
for encoder_layer in self.layers:
layer_outputs = encoder_layer(

6
python/sglang/srt/models/paddleocr_vl.py
View File

@@ -36,7 +36,7 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInp
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.ernie4 import Ernie4_5_ForCausalLM
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
class Projector(nn.Module):
@@ -442,7 +442,9 @@ class SiglipEncoder(nn.Module):
rope_emb = rope_emb_max_grid[pids].flatten(1)
rope_emb = rope_emb.repeat(1, 2)
rope_emb = (rope_emb.cos(), rope_emb.sin())
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu() and isinstance(cu_seqlens, torch.Tensor):
cu_seqlens = cu_seqlens.to("cpu")
attn_cu_seqlens = cu_seqlens
hidden_states = inputs_embeds

6
python/sglang/srt/models/qwen2_5_vl.py
View File

@@ -76,7 +76,7 @@ from sglang.srt.models.utils import RotaryPosMixin, WeightsMapper, permute_inv
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__)
@@ -453,7 +453,9 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
]
)
cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu():
cu_seqlens = cu_seqlens.to("cpu")
# transformers
x = x.unsqueeze(1)
for layer_num, blk in enumerate(self.blocks):

5
python/sglang/srt/models/qwen2_vl.py
View File

@@ -48,7 +48,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.qwen2 import Qwen2Model
from sglang.srt.models.utils import WeightsMapper, compute_cu_seqlens_from_grid_numpy
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
from sglang.srt.utils.hf_transformers_utils import get_processor
logger = logging.getLogger(__name__)
@@ -392,6 +392,9 @@ class Qwen2VisionTransformer(nn.Module):
position_embeddings = (emb.cos(), emb.sin())
# compute cu_seqlens
cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu():
cu_seqlens = cu_seqlens.to("cpu")
# transformers
x = x.unsqueeze(1)

5
python/sglang/srt/models/qwen3_omni_moe.py
View File

@@ -42,7 +42,7 @@ from sglang.srt.models.qwen3_vl_moe import (
Qwen3VLMoeForConditionalGeneration,
load_fused_expert_weights,
)
from sglang.srt.utils import add_prefix, logger
from sglang.srt.utils import add_prefix, is_npu, logger
class Qwen3OmniMoeAudioEncoderLayer(nn.Module):
@@ -278,6 +278,9 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel):
cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(
-1, dtype=torch.int32
)
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if is_npu():
cu_seqlens = cu_seqlens.to("cpu")
for encoder_layer in self.layers:
layer_outputs = encoder_layer(

9
python/sglang/srt/models/qwen3_vl.py
View File

@@ -63,7 +63,7 @@ from sglang.srt.models.utils import (
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix, get_int_env_var
from sglang.srt.utils import add_prefix, get_int_env_var, is_npu
from sglang.srt.utils.hf_transformers_utils import get_processor
logger = logging.getLogger(__name__)
@@ -446,9 +446,12 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
# compute cu_seqlens
cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if not is_npu():
cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
else:
cu_seqlens = cu_seqlens.to("cpu")
x = x.unsqueeze(1)
cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
deepstack_feature_lists = []
num_deepstack_captured = 0