From fce5c75a308413b521f2aa0edac9dde83494eb26 Mon Sep 17 00:00:00 2001 From: Jinyan Yi Date: Mon, 20 Jul 2026 16:20:33 +0800 Subject: [PATCH] [NPU] Fix vit graph tnd cu seqlens (#31701) --- .../npu/graph_runner/vit_npu_graph_runner.py | 5 +++-- python/sglang/srt/managers/mm_utils.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py index b3a909b22..3738c0d36 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py @@ -164,8 +164,9 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner): self.sin_cos_ws[graph_key] = (rotary_pos_emb_cos, rotary_pos_emb_sin) if graph_key not in self.cu_seq_lens: - seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] - self.cu_seq_lens[graph_key] = seq_lens.to("cpu").to(torch.int32) + # TND fused attention expects cumulative seqlens (cu_seqlens[1:]), + # not per-sequence lengths. + self.cu_seq_lens[graph_key] = cu_seqlens[1:].to("cpu").to(torch.int32) if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None: self._create_graph( diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index dd4a6bd74..891d8df6c 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -800,7 +800,7 @@ def _get_chunked_prefill_embedding( is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req) if is_per_image: - if _is_hip: + if _is_hip or _is_npu: # ROCm CI regressed with one large cross-request ViT batch; keep # the previous per-request path on HIP while CUDA uses batching. chunk = _get_chunked_embedding_by_item(