chore: import upstream snapshot with attribution

python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py

@@ -0,0 +1,604 @@
+import re
+from typing import TYPE_CHECKING
+
+import torch
+import torch_npu
+from sgl_kernel_npu.norm.fused_split_qk_norm import fused_split_qk_norm
+
+from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
+    NPUFusedMLAPreprocess,
+    is_fia_nz,
+    is_mla_preprocess_enabled,
+)
+from sglang.srt.layers.attention.dsa.dsa_indexer import scattered_to_tp_attn_full
+from sglang.srt.layers.attention.dsa.utils import (
+    dsa_use_prefill_cp,
+)
+from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+    from sglang.srt.utils import BumpAllocator
+_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
+
+
+# region MHA
+def forward_mha_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+):
+    if m.q_lora_rank is not None:
+        q, latent_cache = (
+            get_attn_tp_context()
+            .fetch_qkv_latent()
+            .split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
+                dim=-1,
+            )
+        )
+
+        # DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
+
+        if m.use_dsa:
+            q_lora = m.q_a_layernorm(q)
+            q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+            _ = m.indexer(
+                x=hidden_states,
+                q_lora=q_lora,
+                positions=positions,
+                forward_batch=forward_batch,
+                layer_id=m.layer_id,
+                return_indices=False,
+            )
+
+        else:
+            q = m.q_a_layernorm(q)
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+            q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+    else:
+        q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+        latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
+
+    _, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
+    kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1)
+    latent_cache = latent_cache.unsqueeze(1)
+
+    if m.use_deepseek_yarn_rope:
+        B, S = q.shape[0], 1
+        cos, sin = m.rotary_emb.get_cos_sin_cache(
+            positions, hidden_states.dtype, offsets=None
+        )
+        q_pe = torch_npu.npu_interleave_rope(
+            q_pe.reshape(B, -1, S, m.qk_rope_head_dim),
+            cos,
+            sin,
+        )
+        q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim)
+
+        ckv_cache, k_rope_cache = get_token_to_kv_pool().get_kv_buffer(m.layer_id)
+        _, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim),  # bnsd
+            m.kv_a_layernorm.weight,
+            cos,
+            sin,
+            forward_batch.out_cache_loc.to(torch.int64),
+            k_rope_cache,
+            ckv_cache,
+            k_rope_scale=None,
+            c_kv_scale=None,
+            k_rope_offset=None,
+            c_kv_offset=None,
+            epsilon=m.kv_a_layernorm.variance_epsilon,
+            cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD",
+            is_output_kv=True,
+        )  # adapter NZ
+
+        k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim)
+    else:
+        kv_a = m.kv_a_layernorm(kv_a)
+        k_pe = latent_cache[:, :, m.kv_lora_rank :]
+        if m.rotary_emb is not None:
+            q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+        # this is for model kimi-vl-a3B-instruct
+        get_token_to_kv_pool().set_kv_buffer(
+            m, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
+        )
+
+    q[..., m.qk_nope_head_dim :] = q_pe
+
+    kv = m.kv_b_proj(kv_a)[0]
+    kv = kv.view(-1, m.num_local_heads, m.qk_nope_head_dim + m.v_head_dim)
+    k_nope = kv[..., : m.qk_nope_head_dim]
+    v = kv[..., m.qk_nope_head_dim :]
+
+    k = m._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
+    return q, k, v, forward_batch
+
+
+def forward_mha_core_npu(
+    m: "DeepseekV2AttentionMLA",
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    forward_batch: "ForwardBatch",
+) -> torch.Tensor:
+    attn_output = m.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
+    attn_output = attn_output.reshape(-1, m.num_local_heads * m.v_head_dim)
+    output, _ = m.o_proj(attn_output)
+    return output
+
+
+# endregion
+
+
+# region MLA
+def forward_mla_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+):
+    if is_mla_preprocess_enabled():
+        if not hasattr(m, "mla_preprocess"):
+            m.mla_preprocess = NPUFusedMLAPreprocess(
+                m.fused_qkv_a_proj_with_mqa,
+                m.q_a_layernorm,
+                m.kv_a_layernorm,
+                m.q_b_proj,
+                m.w_kc,
+                m.rotary_emb,
+                m.layer_id,
+                m.num_local_heads,
+                m.qk_nope_head_dim,
+                m.qk_rope_head_dim,
+                m.quant_config,
+            )
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            forward_batch,
+            zero_allocator,
+            positions,
+        ) = m.mla_preprocess.forward(
+            positions, hidden_states, forward_batch, zero_allocator
+        )
+        topk_indices = None
+    else:
+        q_lora = None
+        if m.q_lora_rank is not None:
+            qkv_latent = get_attn_tp_context().fetch_qkv_latent()
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q, latent_cache = qkv_latent.split(
+                    [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
+                    dim=-1,
+                )
+                k_nope = latent_cache[..., : m.kv_lora_rank]
+
+                q = m.q_a_layernorm(q)
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+
+                k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+                k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+            else:
+                if qkv_latent.shape[0] < 65536 and not dsa_use_prefill_cp(
+                    forward_batch
+                ):
+                    q, k_nope, k_pe = fused_split_qk_norm(
+                        qkv_latent,
+                        m.q_a_layernorm,
+                        m.kv_a_layernorm,
+                        m.q_lora_rank,
+                        m.kv_lora_rank,
+                        m.qk_rope_head_dim,
+                        eps=m.q_a_layernorm.variance_epsilon,
+                    )
+                else:
+                    q, latent_cache = qkv_latent.split(
+                        [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
+                        dim=-1,
+                    )
+                    k_nope = latent_cache[..., : m.kv_lora_rank]
+
+                    q = m.q_a_layernorm(q)
+
+                    k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+                    k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+
+            # q_lora needed by indexer
+            if m.use_dsa:
+                q_lora = q
+
+            q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+        else:
+            q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+            latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
+            k_nope = latent_cache[..., : m.kv_lora_rank]
+            k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+            k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+
+        q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
+
+        q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
+
+        q_nope_out = q_nope_out.transpose(0, 1)
+
+        q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+
+        if dsa_use_prefill_cp(forward_batch):
+            # support allgather+rerrange
+            k_nope, k_pe = m.rebuild_cp_kv_cache(
+                latent_cache, forward_batch, k_nope, k_pe
+            )
+        topk_indices = None
+        if q_lora is not None:
+            topk_indices = m.indexer(
+                x=hidden_states,
+                q_lora=q_lora,
+                positions=positions,
+                forward_batch=forward_batch,
+                layer_id=m.layer_id,
+            )
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        forward_batch,
+        zero_allocator,
+        positions,
+        topk_indices,
+    )
+
+
+def forward_mla_core_npu(
+    m: "DeepseekV2AttentionMLA",
+    q_pe: torch.Tensor,
+    k_pe: torch.Tensor,
+    q_nope_out: torch.Tensor,
+    k_nope: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    positions: torch.Tensor,
+    topk_indices: torch.Tensor,
+) -> torch.Tensor:
+    attn_output = m.attn_mqa(
+        q_nope_out,
+        k_nope,
+        k_nope,
+        forward_batch,
+        q_rope=q_pe,
+        k_rope=k_pe,
+        **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
+    )
+
+    attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
+
+    attn_bmm_output = torch.empty(
+        (attn_output.shape[0], m.num_local_heads, m.v_head_dim),
+        dtype=attn_output.dtype,
+        device=attn_output.device,
+    )
+
+    attn_output = attn_output.contiguous()
+    torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
+
+    attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
+    output, _ = m.o_proj(attn_bmm_output)
+
+    return output
+
+
+# endregion
+
+
+# region DSA
+def forward_dsa_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+    prev_topk_indices: torch.Tensor = None,
+):
+    dynamic_scale = None
+    if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode():
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            q_lora,
+            forward_batch,
+            zero_allocator,
+            positions,
+            dynamic_scale,
+        ) = npu_mla_preprocess(
+            m,
+            hidden_states,
+            positions,
+            forward_batch,
+            zero_allocator,
+        )
+    else:
+        fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+        if m.rotary_emb.is_neox_style:
+            q, latent_cache = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            # overlap qk norm
+            q = m.q_a_layernorm(q)
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+            q_lora = q.clone()  # required for topk_indices
+
+            q_event = None
+            if m.alt_stream is not None:
+                m.alt_stream.wait_stream(torch.npu.current_stream())
+                with torch.npu.stream(m.alt_stream):
+                    q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+                    # record q to ensure memory space will not be released
+                    q.record_stream(m.alt_stream)
+                    q_event = m.alt_stream.record_event()
+            else:
+                q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+            k_nope, k_pe = latent_cache.unsqueeze(1).split(
+                [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
+            )
+            k_nope = m.kv_a_layernorm(k_nope)
+            # main stream waits for the completion of the event on the alt stream to ensure data dependency is complete
+            if q_event is not None:
+                torch.npu.current_stream().wait_event(q_event)
+        else:
+            if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp(
+                forward_batch
+            ):
+                q_lora, k_nope, k_pe = fused_split_qk_norm(
+                    fused_qkv_a_proj_out,
+                    m.q_a_layernorm,
+                    m.kv_a_layernorm,
+                    m.q_lora_rank,
+                    m.kv_lora_rank,
+                    m.qk_rope_head_dim,
+                    eps=m.q_a_layernorm.variance_epsilon,
+                )
+            else:
+                q, latent_cache = fused_qkv_a_proj_out.split(
+                    [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+                )
+                # overlap qk norm
+                q = m.q_a_layernorm(q)
+
+                q_lora = q.clone()  # required for topk_indices
+                k_nope, k_pe = latent_cache.unsqueeze(1).split(
+                    [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
+                )
+                k_nope = m.kv_a_layernorm(k_nope)
+            q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+        q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
+
+        q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
+
+        q_nope_out = q_nope_out.transpose(0, 1)
+
+        if m.layer_id == 0:
+            m.rotary_emb.sin_cos_cache = m.rotary_emb.cos_sin_cache.index_select(
+                0, positions
+            )
+
+        q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+
+        if dsa_use_prefill_cp(forward_batch):
+            # support allgather+rerrange
+            k_nope, k_pe = m.rebuild_cp_kv_cache(
+                latent_cache, forward_batch, k_nope, k_pe
+            )
+
+    if not m.skip_topk or (m.is_nextn and prev_topk_indices is None):
+        topk_indices = m.indexer(
+            hidden_states,
+            q_lora,
+            positions,
+            forward_batch,
+            m.layer_id,
+            layer_scatter_modes,
+            dynamic_scale,
+        )
+    else:
+        topk_indices = prev_topk_indices
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        topk_indices,
+        forward_batch,
+        zero_allocator,
+        positions,
+    )
+
+
+def forward_dsa_core_npu(
+    m: "DeepseekV2AttentionMLA",
+    q_pe: torch.Tensor,
+    k_pe: torch.Tensor,
+    q_nope_out: torch.Tensor,
+    k_nope: torch.Tensor,
+    topk_indices: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    positions: torch.Tensor,
+) -> torch.Tensor:
+    attn_output = m.attn_mqa(
+        q_nope_out.contiguous(),
+        k_nope.contiguous(),
+        k_nope.contiguous(),
+        forward_batch,
+        save_kv_cache=True,  # False if forward_batch.forward_mode.is_extend() else True,
+        q_rope=q_pe.contiguous(),
+        k_rope=k_pe.contiguous(),
+        topk_indices=topk_indices,
+    )
+    attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
+
+    attn_bmm_output = torch.empty(
+        (attn_output.shape[0], m.num_local_heads, m.v_head_dim),
+        dtype=attn_output.dtype,
+        device=attn_output.device,
+    )
+
+    if (
+        forward_batch.forward_mode.is_extend()
+        and not forward_batch.forward_mode.is_draft_extend_v2()
+        and not forward_batch.forward_mode.is_target_verify()
+    ):
+        attn_output = attn_output.transpose(0, 1)
+        torch.bmm(
+            attn_output,
+            m.w_vc,
+            out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(
+                0, 1
+            ),
+        )
+    else:
+        attn_output = attn_output.contiguous()
+        torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
+
+    attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
+
+    output, _ = m.o_proj(attn_bmm_output)
+    if not m.next_skip_topk:
+        return output, None
+    else:
+        return output, topk_indices
+
+
+def npu_mla_preprocess(
+    m: "DeepseekV2AttentionMLA",
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+):
+    dynamic_scale = None
+    if not hasattr(m, "mla_preprocess"):
+        m.mla_preprocess = NPUFusedMLAPreprocess(
+            m.fused_qkv_a_proj_with_mqa,
+            m.q_a_layernorm,
+            m.kv_a_layernorm,
+            m.q_b_proj,
+            m.w_kc,
+            m.rotary_emb,
+            m.layer_id,
+            m.num_local_heads,
+            m.qk_nope_head_dim,
+            m.qk_rope_head_dim,
+            m.v_head_dim,
+            m.quant_config,
+        )
+    # mlaprolog does not require additional calculation of q_lora
+    _is_mlaprolog = hasattr(m.quant_config, "ignore") and any(
+        re.fullmatch(r".*kv_b_proj", l) for l in m.quant_config.ignore
+    )
+    if _is_mlaprolog:
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            q_lora,
+            forward_batch,
+            positions,
+            dynamic_scale,
+        ) = m.mla_preprocess.forward(
+            positions, hidden_states, forward_batch, zero_allocator
+        )
+    else:
+        if m.alt_stream is not None:
+            mla_event = torch.npu.Event()
+            mla_event.record()
+            with torch.npu.stream(m.alt_stream):
+                # alt stream waits for the completion of the event on the main stream to ensure data dependency is complete
+                torch.npu.current_stream().wait_event(mla_event)
+                (
+                    q_pe,
+                    k_pe,
+                    q_nope_out,
+                    k_nope,
+                    forward_batch,
+                    zero_allocator,
+                    positions,
+                ) = m.mla_preprocess.forward(
+                    positions, hidden_states, forward_batch, zero_allocator
+                )
+
+            fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+            q, _ = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            q_lora = m.q_a_layernorm(q)
+            torch.npu.current_stream().wait_event(m.alt_stream)
+        else:
+            (
+                q_pe,
+                k_pe,
+                q_nope_out,
+                k_nope,
+                forward_batch,
+                zero_allocator,
+                positions,
+            ) = m.mla_preprocess.forward(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+            fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+            q, _ = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            q_lora = m.q_a_layernorm(q)
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        q_lora,
+        forward_batch,
+        zero_allocator,
+        positions,
+        dynamic_scale,
+    )
+
+
+# endregion

python/sglang/srt/hardware_backend/npu/modules/glm46v_processor.py

@@ -0,0 +1,285 @@
+"""NPU patch for GLM-4.6V image and video preprocessing.
+
+The GLM-4.6V image processor (Glm46VImageProcessorFast) and video processor
+(Glm46VVideoProcessor) create 10-dimensional tensors during patch extraction,
+which exceeds Ascend NPU's 8-dimension limit.
+
+This patch restructures the computation to stay within 8 dimensions, following
+the same pattern as the Qwen VL NPU patch.
+"""
+
+import torch
+import torchvision.transforms.v2.functional as tvF
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_processing_utils_fast import (
+    group_images_by_shape,
+    reorder_images,
+)
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from transformers.models.glm46v.image_processing_glm46v import smart_resize
+from transformers.utils import TensorType
+from transformers.video_utils import group_videos_by_shape, reorder_videos
+
+from sglang.srt.hardware_backend.npu.modules.qwen_vl_processor import (
+    transform_patches_to_flatten,
+)
+from sglang.srt.utils import apply_module_patch
+
+
+# Func refers to transformers.models.glm46v.image_processing_glm46v_fast.py
+# Glm46VImageProcessorFast._preprocess
+def npu_wrapper_glm46v_preprocess(func):
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ):
+        grouped_images, grouped_images_index = group_images_by_shape(
+            images, disable_grouping=disable_grouping
+        )
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            height, width = stacked_images.shape[-2:]
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    num_frames=temporal_patch_size,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_images = self.resize(
+                    stacked_images,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+            resized_images_grouped[shape] = stacked_images
+
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        grouped_images, grouped_images_index = group_images_by_shape(
+            resized_images, disable_grouping=disable_grouping
+        )
+        processed_images_grouped = {}
+        processed_grids = {}
+
+        for shape, stacked_images in grouped_images.items():
+            resized_height, resized_width = stacked_images.shape[-2:]
+
+            patches = self.rescale_and_normalize(
+                stacked_images,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            if patches.ndim == 4:
+                patches = patches.unsqueeze(1)
+
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(
+                    1,
+                    temporal_patch_size - (patches.shape[1] % temporal_patch_size),
+                    1,
+                    1,
+                    1,
+                )
+                patches = torch.cat([patches, repeats], dim=1)
+
+            batch_size, t_len, channel = patches.shape[:3]
+            grid_t = t_len // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_images_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_images = reorder_images(
+            processed_images_grouped, grouped_images_index
+        )
+        processed_grids = reorder_images(processed_grids, grouped_images_index)
+
+        pixel_values = torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    return _preprocess
+
+
+# Func refers to transformers.models.glm46v.video_processing_glm46v.py
+# Glm46VVideoProcessor._preprocess
+def npu_wrapper_glm46v_video_preprocess(func):
+
+    def _preprocess(
+        self,
+        videos: list[torch.Tensor],
+        do_convert_rgb: bool = True,
+        do_resize: bool = True,
+        size: SizeDict | None = None,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+
+        for shape, stacked_videos in grouped_videos.items():
+            B, T, C, H, W = stacked_videos.shape
+            num_frames, height, width = T, H, W
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_videos = stacked_videos.view(B * T, C, H, W)
+                stacked_videos = self.resize(
+                    stacked_videos,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+                stacked_videos = stacked_videos.view(
+                    B, T, C, resized_height, resized_width
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            resized_height, resized_width = get_image_size(
+                stacked_videos[0], channel_dim=ChannelDimension.FIRST
+            )
+
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            patches = stacked_videos
+
+            # Check that videos have `num_frames` divisible by `temporal_patch_size`
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+                patches = torch.cat([patches, repeats], dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_videos_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_videos = reorder_videos(
+            processed_videos_grouped, grouped_videos_index
+        )
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+        data = {
+            "pixel_values_videos": pixel_values_videos,
+            "video_grid_thw": video_grid_thw,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    return _preprocess
+
+
+_npu_glm46v_preprocess_patched = False
+
+
+def npu_apply_glm46v_image_preprocess_patch():
+    global _npu_glm46v_preprocess_patched
+    if _npu_glm46v_preprocess_patched:
+        return
+    apply_module_patch(
+        "transformers.models.glm46v.image_processing_glm46v_fast.Glm46VImageProcessorFast",
+        "_preprocess",
+        [npu_wrapper_glm46v_preprocess],
+    )
+    apply_module_patch(
+        "transformers.models.glm46v.video_processing_glm46v.Glm46VVideoProcessor",
+        "_preprocess",
+        [npu_wrapper_glm46v_video_preprocess],
+    )
+    _npu_glm46v_preprocess_patched = True

python/sglang/srt/hardware_backend/npu/modules/qwen_vl_processor.py

@@ -0,0 +1,304 @@
+import torch
+import torchvision.transforms.v2.functional as tvF
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_transforms import group_images_by_shape, reorder_images
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
+    smart_resize as smart_resize_video,
+)
+from transformers.utils import TensorType
+from transformers.video_utils import group_videos_by_shape, reorder_videos
+
+from sglang.srt.utils import apply_module_patch
+
+
+def transform_patches_to_flatten(
+    patches: torch.Tensor,
+    batch_size: int,
+    grid_t: int,
+    temporal_patch_size: int,
+    channel: int,
+    grid_h: int,
+    grid_w: int,
+    patch_size: int,
+    merge_size: int,
+) -> torch.Tensor:
+    patches = patches.view(
+        batch_size * grid_t,
+        temporal_patch_size * channel,
+        grid_h // merge_size,
+        merge_size,
+        patch_size,
+        grid_w // merge_size,
+        merge_size,
+        patch_size,
+    )
+    patches = patches.permute(0, 1, 2, 5, 3, 6, 4, 7)
+    patches = patches.reshape(
+        batch_size,
+        grid_t,
+        temporal_patch_size,
+        channel,
+        grid_h * grid_w,
+        patch_size,
+        patch_size,
+    )
+    patches = patches.permute(0, 1, 4, 3, 2, 5, 6)
+    flatten_patches = patches.reshape(
+        batch_size,
+        grid_t * grid_h * grid_w,
+        -1,
+    )
+    return flatten_patches
+
+
+# Func refers to transformers.models.qwen2_vl.image_processing_qwen2_vl.py
+# Qwen2VLImageProcessor._preprocess
+def npu_wrapper_preprocess(func):
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ):
+        # Group images by size for batched resizing
+        grouped_images, grouped_images_index = group_images_by_shape(
+            images, disable_grouping=disable_grouping
+        )
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            height, width = stacked_images.shape[-2:]
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    height,
+                    width,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_images = self.resize(
+                    image=stacked_images,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+            resized_images_grouped[shape] = stacked_images
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        # Group images by size for further processing
+        # Needed in case do_resize is False, or resize returns images with different sizes
+        grouped_images, grouped_images_index = group_images_by_shape(
+            resized_images, disable_grouping=disable_grouping
+        )
+        processed_images_grouped = {}
+        processed_grids = {}
+        for shape, stacked_images in grouped_images.items():
+            resized_height, resized_width = stacked_images.shape[-2:]
+            # Fused rescale and normalize
+            patches = self.rescale_and_normalize(
+                stacked_images,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            if patches.ndim == 4:
+                # add a temporal dimension if we have images
+                patches = patches.unsqueeze(1)
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+                patches = torch.cat([patches, repeats], dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_images_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_images = reorder_images(
+            processed_images_grouped, grouped_images_index
+        )
+        processed_grids = reorder_images(processed_grids, grouped_images_index)
+        pixel_values = torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    return _preprocess
+
+
+# Func refers to transformers.models.qwen3_vl.video_processing_qwen3_vl.py
+# Qwen3VLVideoProcessor._preprocess
+def npu_wrapper_video_preprocess(func):
+
+    def _preprocess(
+        self,
+        videos: list[torch.Tensor],
+        do_convert_rgb: bool = True,
+        do_resize: bool = True,
+        size: SizeDict | None = None,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+
+        for shape, stacked_videos in grouped_videos.items():
+            B, T, C, H, W = stacked_videos.shape
+            num_frames, height, width = T, H, W
+            if do_resize:
+                resized_height, resized_width = smart_resize_video(
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_videos = stacked_videos.view(B * T, C, H, W)
+                stacked_videos = self.resize(
+                    stacked_videos,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+                stacked_videos = stacked_videos.view(
+                    B, T, C, resized_height, resized_width
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            resized_height, resized_width = get_image_size(
+                stacked_videos[0], channel_dim=ChannelDimension.FIRST
+            )
+
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            patches = stacked_videos
+
+            # Check that videos have `num_frames` divisible by `temporal_patch_size`
+            T = patches.shape[1]
+            if pad := -T % temporal_patch_size:
+                repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)
+                patches = torch.cat((patches, repeats), dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_videos_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_videos = reorder_videos(
+            processed_videos_grouped, grouped_videos_index
+        )
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+        data = {
+            "pixel_values_videos": pixel_values_videos,
+            "video_grid_thw": video_grid_thw,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    return _preprocess
+
+
+_npu_preprocess_patched = False
+
+
+def npu_apply_qwen_image_preprocess_patch():
+    global _npu_preprocess_patched
+    if _npu_preprocess_patched:
+        return
+    apply_module_patch(
+        "transformers.models.qwen2_vl.image_processing_qwen2_vl.Qwen2VLImageProcessor",
+        "_preprocess",
+        [npu_wrapper_preprocess],
+    )
+    apply_module_patch(
+        "transformers.models.qwen3_vl.video_processing_qwen3_vl.Qwen3VLVideoProcessor",
+        "_preprocess",
+        [npu_wrapper_video_preprocess],
+    )
+    _npu_preprocess_patched = True