diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 001934ed6..ff603e8f8 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -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: diff --git a/python/sglang/srt/models/dots_vlm_vit.py b/python/sglang/srt/models/dots_vlm_vit.py index ca82ddd5e..873994e0b 100644 --- a/python/sglang/srt/models/dots_vlm_vit.py +++ b/python/sglang/srt/models/dots_vlm_vit.py @@ -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( diff --git a/python/sglang/srt/models/glm4v.py b/python/sglang/srt/models/glm4v.py index cee2dd8d3..243677f82 100644 --- a/python/sglang/srt/models/glm4v.py +++ b/python/sglang/srt/models/glm4v.py @@ -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) diff --git a/python/sglang/srt/models/idefics2.py b/python/sglang/srt/models/idefics2.py index 02f8a2497..c16c86d10 100644 --- a/python/sglang/srt/models/idefics2.py +++ b/python/sglang/srt/models/idefics2.py @@ -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( diff --git a/python/sglang/srt/models/paddleocr_vl.py b/python/sglang/srt/models/paddleocr_vl.py index 89ede652b..456fb19ab 100644 --- a/python/sglang/srt/models/paddleocr_vl.py +++ b/python/sglang/srt/models/paddleocr_vl.py @@ -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 diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index aafde99bb..7f44978fc 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/python/sglang/srt/models/qwen2_5_vl.py @@ -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): diff --git a/python/sglang/srt/models/qwen2_vl.py b/python/sglang/srt/models/qwen2_vl.py index 24ec4c430..3002bbe14 100644 --- a/python/sglang/srt/models/qwen2_vl.py +++ b/python/sglang/srt/models/qwen2_vl.py @@ -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) diff --git a/python/sglang/srt/models/qwen3_omni_moe.py b/python/sglang/srt/models/qwen3_omni_moe.py index 045ba2eba..5c80458b6 100644 --- a/python/sglang/srt/models/qwen3_omni_moe.py +++ b/python/sglang/srt/models/qwen3_omni_moe.py @@ -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( diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 891913078..6c06aaa0e 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -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