sglang 0.4.9__py3-none-any.whl → 0.4.9.post2__py3-none-any.whl

sglang/srt/mem_cache/memory_pool.py

@@ -37,12 +37,15 @@ import triton.language as tl
 
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.layers.radix_attention import RadixAttention
-from sglang.srt.utils import debug_timing, get_bool_env_var, is_cuda, next_power_of_2
+from sglang.srt.utils import get_bool_env_var, is_cuda, is_npu, next_power_of_2
 
 logger = logging.getLogger(__name__)
 
 GB = 1024 * 1024 * 1024
 _is_cuda = is_cuda()
+_is_npu = is_npu()
+if not _is_npu:
+    from sgl_kernel.kvcacheio import transfer_kv_per_layer, transfer_kv_per_layer_mla
 
 
 class ReqToTokenPool:
@@ -150,13 +153,16 @@ class KVCache(abc.ABC):
     ) -> None:
         raise NotImplementedError()
 
-    def get_flat_data(self, indices):
-        raise NotImplementedError()
-
-    def transfer(self, indices, flat_data):
+    @abc.abstractmethod
+    def load_from_host_per_layer(
+        self, host_pool, host_indices, device_indices, layer_id, io_backend
+    ):
         raise NotImplementedError()
 
-    def transfer_per_layer(self, indices, flat_data, layer_id):
+    @abc.abstractmethod
+    def backup_to_host_all_layer(
+        self, host_pool, host_indices, device_indices, io_backend
+    ):
         raise NotImplementedError()
 
     def register_layer_transfer_counter(self, layer_transfer_counter):
@@ -247,7 +253,7 @@ class MHATokenToKVPool(KVCache):
                     )
                     for _ in range(self.layer_num)
                 ]
-
+        self.token_stride = self.head_num * self.head_dim
         self.data_ptrs = torch.tensor(
             [x.data_ptr() for x in self.k_buffer + self.v_buffer],
             dtype=torch.uint64,
@@ -281,24 +287,24 @@ class MHATokenToKVPool(KVCache):
         # layer_num x [seq_len, head_num, head_dim]
         # layer_num x [page_num, page_size, head_num, head_dim]
         kv_data_ptrs = [
-            self.get_key_buffer(i).data_ptr()
+            self._get_key_buffer(i).data_ptr()
             for i in range(self.start_layer, self.start_layer + self.layer_num)
         ] + [
-            self.get_value_buffer(i).data_ptr()
+            self._get_value_buffer(i).data_ptr()
             for i in range(self.start_layer, self.start_layer + self.layer_num)
         ]
         kv_data_lens = [
-            self.get_key_buffer(i).nbytes
+            self._get_key_buffer(i).nbytes
             for i in range(self.start_layer, self.start_layer + self.layer_num)
         ] + [
-            self.get_value_buffer(i).nbytes
+            self._get_value_buffer(i).nbytes
             for i in range(self.start_layer, self.start_layer + self.layer_num)
         ]
         kv_item_lens = [
-            self.get_key_buffer(i)[0].nbytes * self.page_size
+            self._get_key_buffer(i)[0].nbytes * self.page_size
             for i in range(self.start_layer, self.start_layer + self.layer_num)
         ] + [
-            self.get_value_buffer(i)[0].nbytes * self.page_size
+            self._get_value_buffer(i)[0].nbytes * self.page_size
             for i in range(self.start_layer, self.start_layer + self.layer_num)
         ]
         return kv_data_ptrs, kv_data_lens, kv_item_lens
@@ -341,49 +347,73 @@ class MHATokenToKVPool(KVCache):
                 self.v_buffer[layer_id][chunk_indices] = v_chunk
         torch.cuda.synchronize()
 
-    # Todo: different memory layout
-    def get_flat_data(self, indices):
-        # prepare a large chunk of contiguous data for efficient transfer
-        flatten = torch.stack(
-            [
-                torch.stack([self.k_buffer[i][indices] for i in range(self.layer_num)]),
-                torch.stack([self.v_buffer[i][indices] for i in range(self.layer_num)]),
-            ]
+    def load_from_host_per_layer(
+        self,
+        host_pool,
+        host_indices,
+        device_indices,
+        layer_id,
+        io_backend,
+    ):
+        transfer_kv_per_layer(
+            src_k=host_pool.k_buffer[layer_id],
+            dst_k=self.k_buffer[layer_id],
+            src_v=host_pool.v_buffer[layer_id],
+            dst_v=self.v_buffer[layer_id],
+            src_indices=host_indices,
+            dst_indices=device_indices,
+            io_backend=io_backend,
+            page_size=self.page_size,
+            item_size=self.token_stride,
         )
-        return flatten
-
-    @debug_timing
-    def transfer(self, indices, flat_data):
-        # transfer prepared data from host to device
-        flat_data = flat_data.to(device=self.device, non_blocking=False)
-        k_data, v_data = flat_data[0], flat_data[1]
-        for i in range(self.layer_num):
-            self.k_buffer[i][indices] = k_data[i]
-            self.v_buffer[i][indices] = v_data[i]
-
-    def transfer_per_layer(self, indices, flat_data, layer_id):
-        # transfer prepared data from host to device
-        flat_data = flat_data.to(device=self.device, non_blocking=False)
-        k_data, v_data = flat_data[0], flat_data[1]
-        self.k_buffer[layer_id - self.start_layer][indices] = k_data
-        self.v_buffer[layer_id - self.start_layer][indices] = v_data
 
-    def get_key_buffer(self, layer_id: int):
-        if self.layer_transfer_counter is not None:
-            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+    def backup_to_host_all_layer(
+        self, host_pool, host_indices, device_indices, io_backend
+    ):
+        # todo: specialized all layer kernels for the layer-non-contiguous memory pool
+        for layer_id in range(self.start_layer, self.start_layer + self.layer_num):
+            if layer_id - self.start_layer >= len(host_pool.k_buffer):
+                raise ValueError(
+                    f"Layer ID {layer_id} exceeds the number of layers in host pool."
+                )
+            transfer_kv_per_layer(
+                src_k=self.k_buffer[layer_id],
+                dst_k=host_pool.k_buffer[layer_id],
+                src_v=self.v_buffer[layer_id],
+                dst_v=host_pool.v_buffer[layer_id],
+                src_indices=device_indices,
+                dst_indices=host_indices,
+                io_backend=io_backend,
+                page_size=self.page_size,
+                item_size=self.token_stride,
+            )
 
+    def _get_key_buffer(self, layer_id: int):
+        # for internal use of referencing
         if self.store_dtype != self.dtype:
             return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
         return self.k_buffer[layer_id - self.start_layer]
 
-    def get_value_buffer(self, layer_id: int):
+    def get_key_buffer(self, layer_id: int):
+        # note: get_key_buffer is hooked with synchronization for layer-wise KV cache loading
+        # it is supposed to be used only by attention backend not for information purpose
+        # same applies to get_value_buffer and get_kv_buffer
         if self.layer_transfer_counter is not None:
             self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
 
+        return self._get_key_buffer(layer_id)
+
+    def _get_value_buffer(self, layer_id: int):
+        # for internal use of referencing
         if self.store_dtype != self.dtype:
             return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
         return self.v_buffer[layer_id - self.start_layer]
 
+    def get_value_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+        return self._get_value_buffer(layer_id)
+
     def get_kv_buffer(self, layer_id: int):
         return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id)
 
@@ -574,32 +604,49 @@ class AscendTokenToKVPool(MHATokenToKVPool):
         with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
             # [size, head_num, head_dim] for each layer
             # The padded slot 0 is used for writing dummy outputs from padded tokens.
-            self.k_buffer = [
-                torch.zeros(
-                    (
-                        self.size // self.page_size + 1,
-                        self.page_size,
-                        self.head_num,
-                        self.head_dim,
-                    ),
-                    dtype=self.store_dtype,
-                    device=self.device,
-                )
-                for _ in range(self.layer_num)
-            ]
-            self.v_buffer = [
-                torch.zeros(
-                    (
-                        self.size // self.page_size + 1,
-                        self.page_size,
-                        self.head_num,
-                        self.head_dim,
-                    ),
-                    dtype=self.store_dtype,
-                    device=self.device,
-                )
-                for _ in range(self.layer_num)
-            ]
+            # Continuous memory improves the efficiency of Ascend`s transmission backend,
+            # while other backends remain unchanged.
+            self.kv_buffer = torch.zeros(
+                (
+                    2,
+                    self.layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    self.head_num,
+                    self.head_dim,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+            self.k_buffer = self.kv_buffer[0]
+            self.v_buffer = self.kv_buffer[1]
+
+    # for disagg
+    def get_contiguous_buf_infos(self):
+        # layer_num x [seq_len, head_num, head_dim]
+        # layer_num x [page_num, page_size, head_num, head_dim]
+        kv_data_ptrs = [
+            self.get_key_buffer(i).data_ptr()
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ] + [
+            self.get_value_buffer(i).data_ptr()
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ]
+        kv_data_lens = [
+            self.get_key_buffer(i).nbytes
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ] + [
+            self.get_value_buffer(i).nbytes
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ]
+        kv_item_lens = [
+            self.get_key_buffer(i)[0].nbytes
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ] + [
+            self.get_value_buffer(i)[0].nbytes
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ]
+        return kv_data_ptrs, kv_data_lens, kv_item_lens
 
     def set_kv_buffer(
         self,
@@ -761,6 +808,7 @@ class MLATokenToKVPool(KVCache):
                     for _ in range(layer_num)
                 ]
 
+        self.token_stride = kv_lora_rank + qk_rope_head_dim
         self.layer_transfer_counter = None
 
         kv_size = self.get_kv_size_bytes()
@@ -846,21 +894,37 @@ class MLATokenToKVPool(KVCache):
             self.kv_buffer[layer_id], loc, cache_k_nope, cache_k_rope
         )
 
-    def get_flat_data(self, indices):
-        # prepare a large chunk of contiguous data for efficient transfer
-        return torch.stack([self.kv_buffer[i][indices] for i in range(self.layer_num)])
-
-    @debug_timing
-    def transfer(self, indices, flat_data):
-        # transfer prepared data from host to device
-        flat_data = flat_data.to(device=self.device, non_blocking=False)
-        for i in range(self.layer_num):
-            self.kv_buffer[i][indices] = flat_data[i]
+    def load_from_host_per_layer(
+        self, host_pool, host_indices, device_indices, layer_id, io_backend
+    ):
+        transfer_kv_per_layer_mla(
+            src=host_pool.kv_buffer[layer_id],
+            dst=self.kv_buffer[layer_id],
+            src_indices=host_indices,
+            dst_indices=device_indices,
+            io_backend=io_backend,
+            page_size=self.page_size,
+            item_size=self.token_stride,
+        )
 
-    def transfer_per_layer(self, indices, flat_data, layer_id):
-        # transfer prepared data from host to device
-        flat_data = flat_data.to(device=self.device, non_blocking=False)
-        self.kv_buffer[layer_id - self.start_layer][indices] = flat_data
+    def backup_to_host_all_layer(
+        self, host_pool, host_indices, device_indices, io_backend
+    ):
+        # todo: specialized all layer kernels for the layer-non-contiguous memory pool
+        for layer_id in range(self.start_layer, self.start_layer + self.layer_num):
+            if layer_id - self.start_layer >= len(host_pool.kv_buffer):
+                raise ValueError(
+                    f"Layer ID {layer_id} exceeds the number of layers in host pool."
+                )
+            transfer_kv_per_layer_mla(
+                src=self.kv_buffer[layer_id],
+                dst=host_pool.kv_buffer[layer_id],
+                src_indices=device_indices,
+                dst_indices=host_indices,
+                io_backend=io_backend,
+                page_size=self.page_size,
+                item_size=self.token_stride,
+            )
 
     def get_cpu_copy(self, indices):
         torch.cuda.synchronize()
@@ -922,18 +986,16 @@ class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
 
         with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
             # The padded slot 0 is used for writing dummy outputs from padded tokens.
-            self.kv_buffer = [
-                torch.zeros(
-                    (
-                        self.size // self.page_size + 1,
-                        self.page_size,
-                        self.kv_lora_rank + self.qk_rope_head_dim,
-                    ),
-                    dtype=self.store_dtype,
-                    device=self.device,
-                )
-                for _ in range(layer_num)
-            ]
+            self.kv_buffer = torch.zeros(
+                (
+                    layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    self.kv_lora_rank + self.qk_rope_head_dim,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
 
         self.layer_transfer_counter = None
 
@@ -943,6 +1005,14 @@ class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
         )
         self.mem_usage = kv_size / GB
 
+    # for disagg
+    def get_contiguous_buf_infos(self):
+        # MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
+        kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in range(self.layer_num)]
+        kv_data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
+        kv_item_lens = [self.kv_buffer[i][0].nbytes for i in range(self.layer_num)]
+        return kv_data_ptrs, kv_data_lens, kv_item_lens
+
     def set_kv_buffer(
         self,
         layer: RadixAttention,
@@ -1046,14 +1116,19 @@ class DoubleSparseTokenToKVPool(KVCache):
         self.v_buffer[layer_id - self.start_layer][loc] = cache_v
         self.label_buffer[layer_id - self.start_layer][loc] = cache_label
 
-    def get_flat_data(self, indices):
-        pass
-
-    def transfer(self, indices, flat_data):
-        pass
+    def load_from_host_per_layer(
+        self, host_pool, host_indices, device_indices, layer_id, io_backend
+    ):
+        raise NotImplementedError(
+            "HiCache not supported for DoubleSparseTokenToKVPool."
+        )
 
-    def transfer_per_layer(self, indices, flat_data, layer_id):
-        pass
+    def backup_to_host_all_layer(
+        self, host_pool, host_indices, device_indices, io_backend
+    ):
+        raise NotImplementedError(
+            "HiCache not supported for DoubleSparseTokenToKVPool."
+        )
 
 
 @triton.jit

sglang/srt/mem_cache/memory_pool_host.py

@@ -8,7 +8,6 @@ import psutil
 import torch
 
 from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool, MLATokenToKVPool
-from sglang.srt.utils import debug_timing
 
 logger = logging.getLogger(__name__)
 
@@ -99,22 +98,6 @@ class HostKVCache(abc.ABC):
     def init_kv_buffer(self):
         raise NotImplementedError()
 
-    @abc.abstractmethod
-    def transfer(self, indices, flat_data):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get_flat_data(self, indices):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get_flat_data_by_layer(self, indices, layer_id):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def assign_flat_data(self, indices, flat_data):
-        raise NotImplementedError()
-
     @synchronized()
     def clear(self):
         # Initialize memory states and tracking structures.
@@ -243,58 +226,13 @@ class MHATokenToKVPoolHost(HostKVCache):
             pin_memory=self.pin_memory,
         )
 
-    @debug_timing
-    def transfer(self, indices, flat_data):
-        # backup prepared data from device to host
-        self.kv_buffer[:, :, indices] = flat_data.to(
-            device=self.device, non_blocking=False
-        )
+    @property
+    def k_buffer(self):
+        return self.kv_buffer[0]
 
-    def get_flat_data(self, indices):
-        return self.kv_buffer[:, :, indices]
-
-    def get_flat_data_by_layer(self, indices, layer_id):
-        return self.kv_buffer[:, layer_id - self.start_layer, indices]
-
-    def assign_flat_data(self, indices, flat_data):
-        self.kv_buffer[:, :, indices] = flat_data
-
-    def write_page_all_layers(self, host_indices, device_indices, device_pool):
-        device_indices_cpu = device_indices[:: self.page_size].cpu()
-        for i in range(len(device_indices_cpu)):
-            h_index = host_indices[i * self.page_size]
-            d_index = device_indices_cpu[i]
-            for j in range(self.layer_num):
-                self.kv_buffer[0, j, h_index : h_index + self.page_size].copy_(
-                    device_pool.k_buffer[j][d_index : d_index + self.page_size],
-                    non_blocking=True,
-                )
-                self.kv_buffer[1, j, h_index : h_index + self.page_size].copy_(
-                    device_pool.v_buffer[j][d_index : d_index + self.page_size],
-                    non_blocking=True,
-                )
-
-    def load_page_per_layer(self, host_indices, device_indices, device_pool, layer_id):
-        device_indices_cpu = device_indices[:: self.page_size].cpu()
-        for i in range(len(device_indices_cpu)):
-            h_index = host_indices[i * self.page_size]
-            d_index = device_indices_cpu[i]
-            device_pool.k_buffer[layer_id - self.start_layer][
-                d_index : d_index + self.page_size
-            ].copy_(
-                self.kv_buffer[
-                    0, layer_id - self.start_layer, h_index : h_index + self.page_size
-                ],
-                non_blocking=True,
-            )
-            device_pool.v_buffer[layer_id - self.start_layer][
-                d_index : d_index + self.page_size
-            ].copy_(
-                self.kv_buffer[
-                    1, layer_id - self.start_layer, h_index : h_index + self.page_size
-                ],
-                non_blocking=True,
-            )
+    @property
+    def v_buffer(self):
+        return self.kv_buffer[1]
 
 
 class MLATokenToKVPoolHost(HostKVCache):
@@ -337,44 +275,3 @@ class MLATokenToKVPoolHost(HostKVCache):
             device=self.device,
             pin_memory=self.pin_memory,
         )
-
-    @debug_timing
-    def transfer(self, indices, flat_data):
-        # backup prepared data from device to host
-        self.kv_buffer[:, indices] = flat_data.to(
-            device=self.device, non_blocking=False
-        )
-
-    def get_flat_data(self, indices):
-        return self.kv_buffer[:, indices]
-
-    def get_flat_data_by_layer(self, indices, layer_id):
-        return self.kv_buffer[layer_id - self.start_layer, indices]
-
-    def assign_flat_data(self, indices, flat_data):
-        self.kv_buffer[:, indices] = flat_data
-
-    def write_page_all_layers(self, host_indices, device_indices, device_pool):
-        device_indices_cpu = device_indices[:: self.page_size].cpu()
-        for i in range(len(device_indices_cpu)):
-            h_index = host_indices[i * self.page_size]
-            d_index = device_indices_cpu[i]
-            for j in range(self.layer_num):
-                self.kv_buffer[j, h_index : h_index + self.page_size].copy_(
-                    device_pool.kv_buffer[j][d_index : d_index + self.page_size],
-                    non_blocking=True,
-                )
-
-    def load_page_per_layer(self, host_indices, device_indices, device_pool, layer_id):
-        device_indices_cpu = device_indices[:: self.page_size].cpu()
-        for i in range(len(device_indices_cpu)):
-            h_index = host_indices[i * self.page_size]
-            d_index = device_indices_cpu[i]
-            device_pool.kv_buffer[layer_id - self.start_layer][
-                d_index : d_index + self.page_size
-            ].copy_(
-                self.kv_buffer[
-                    layer_id - self.start_layer, h_index : h_index + self.page_size
-                ],
-                non_blocking=True,
-            )

sglang/srt/mem_cache/radix_cache.py

@@ -196,11 +196,13 @@ class RadixCache(BasePrefixCache):
 
         if self.page_size != 1:
             page_aligned_len = len(kv_indices) // self.page_size * self.page_size
-            page_aligned_kv_indices = kv_indices[:page_aligned_len].clone()
+            page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
+                dtype=torch.int64, copy=True
+            )
             self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:])
         else:
             page_aligned_len = len(kv_indices)
-            page_aligned_kv_indices = kv_indices.clone()
+            page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
 
         # Radix Cache takes one ref in memory pool
         new_prefix_len = self.insert(
@@ -226,10 +228,12 @@ class RadixCache(BasePrefixCache):
 
         if self.page_size != 1:
             page_aligned_len = len(kv_indices) // self.page_size * self.page_size
-            page_aligned_kv_indices = kv_indices[:page_aligned_len].clone()
+            page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
+                dtype=torch.int64, copy=True
+            )
         else:
             page_aligned_len = len(kv_indices)
-            page_aligned_kv_indices = kv_indices.clone()
+            page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
         page_aligned_token_ids = token_ids[:page_aligned_len]
 
         # Radix Cache takes one ref in memory pool

sglang/srt/model_executor/forward_batch_info.py

@@ -453,8 +453,20 @@ class ForwardBatch:
             for mm_input in self.mm_inputs
         )
 
+    def contains_video_inputs(self) -> bool:
+        if self.mm_inputs is None:
+            return False
+        return any(
+            mm_input is not None and mm_input.contains_video_inputs()
+            for mm_input in self.mm_inputs
+        )
+
     def contains_mm_inputs(self) -> bool:
-        return self.contains_audio_inputs() or self.contains_image_inputs()
+        return (
+            self.contains_audio_inputs()
+            or self.contains_video_inputs()
+            or self.contains_image_inputs()
+        )
 
     def _compute_mrope_positions(
         self, model_runner: ModelRunner, batch: ModelWorkerBatch

sglang/srt/model_loader/loader.py

@@ -64,10 +64,13 @@ from sglang.srt.model_loader.weight_utils import (
 from sglang.srt.utils import (
     get_bool_env_var,
     get_device_capability,
+    is_npu,
     is_pin_memory_available,
     set_weight_attrs,
 )
 
+_is_npu = is_npu()
+
 
 @contextmanager
 def device_loading_context(module: torch.nn.Module, target_device: torch.device):
@@ -127,18 +130,19 @@ def _get_quantization_config(
         # (yizhang2077) workaround for nvidia/Llama-4-Maverick-17B-128E-Eagle3
         if quant_config is None:
             return None
-        major, minor = get_device_capability()
-
-        if major is not None and minor is not None:
-            assert 0 <= minor < 10
-            capability = major * 10 + minor
-            if capability < quant_config.get_min_capability():
-                raise ValueError(
-                    f"The quantization method {model_config.quantization} "
-                    "is not supported for the current GPU. "
-                    f"Minimum capability: {quant_config.get_min_capability()}. "
-                    f"Current capability: {capability}."
-                )
+        if not _is_npu:
+            major, minor = get_device_capability()
+
+            if major is not None and minor is not None:
+                assert 0 <= minor < 10
+                capability = major * 10 + minor
+                if capability < quant_config.get_min_capability():
+                    raise ValueError(
+                        f"The quantization method {model_config.quantization} "
+                        "is not supported for the current GPU. "
+                        f"Minimum capability: {quant_config.get_min_capability()}. "
+                        f"Current capability: {capability}."
+                    )
         supported_dtypes = quant_config.get_supported_act_dtypes()
         if model_config.dtype not in supported_dtypes:
             raise ValueError(
@@ -157,6 +161,13 @@ def _initialize_model(
     """Initialize a model with the given configurations."""
     model_class, _ = get_model_architecture(model_config)
     packed_modules_mapping = getattr(model_class, "packed_modules_mapping", {})
+    if _is_npu:
+        packed_modules_mapping["fused_qkv_a_proj_with_mqa"] = [
+            "q_a_proj",
+            "kv_a_proj_with_mqa",
+        ]
+        packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]
+        packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
     quant_config = _get_quantization_config(
         model_config, load_config, packed_modules_mapping
     )

sglang/srt/models/deepseek_janus_pro.py

@@ -1989,7 +1989,7 @@ class MultiModalityCausalLM(MultiModalityPreTrainedModel):
         hidden_states = general_mm_embed_routine(
             input_ids=input_ids,
             forward_batch=forward_batch,
-            image_data_embedding_func=self.get_image_feature,
+            multimodal_model=self,
             language_model=self.language_model,
             positions=positions,
         )