sglang 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

sglang/srt/managers/{router/manager.py → controller/manager_single.py}

@@ -1,20 +1,26 @@
+"""A controller that manages a group of tensor parallel workers."""
+
 import asyncio
 import logging
+from concurrent.futures import ThreadPoolExecutor
 
 import uvloop
 import zmq
 import zmq.asyncio
 
 from sglang.global_config import global_config
-from sglang.srt.managers.router.model_rpc import ModelRpcClient
+from sglang.srt.managers.controller.tp_worker import ModelTpClient
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback
+from sglang.srt.utils import kill_parent_process
+from sglang.utils import get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
+logger = logging.getLogger("srt.controller")
+
 
-class RouterManager:
-    def __init__(self, model_client: ModelRpcClient, port_args: PortArgs):
+class ControllerSingle:
+    def __init__(self, model_client: ModelTpClient, port_args: PortArgs):
         # Init communication
         context = zmq.asyncio.Context(2)
         self.recv_from_tokenizer = context.socket(zmq.PULL)
@@ -30,7 +36,7 @@ class RouterManager:
         self.recv_reqs = []
 
         # Init some configs
-        self.request_dependency_time = global_config.request_dependency_time
+        self.request_dependency_delay = global_config.request_dependency_delay
 
     async def loop_for_forward(self):
         while True:
@@ -44,14 +50,16 @@ class RouterManager:
             # async sleep for receiving the subsequent request and avoiding cache miss
             slept = False
             if len(out_pyobjs) != 0:
-                has_finished = any([obj.finished for obj in out_pyobjs])
+                has_finished = any(
+                    [obj.finished_reason is not None for obj in out_pyobjs]
+                )
                 if has_finished:
-                    if self.request_dependency_time > 0:
+                    if self.request_dependency_delay > 0:
                         slept = True
-                        await asyncio.sleep(self.request_dependency_time)
+                        await asyncio.sleep(self.request_dependency_delay)
 
             if not slept:
-                await asyncio.sleep(0.0006)
+                await asyncio.sleep(global_config.wait_for_new_request_delay)
 
     async def loop_for_recv_requests(self):
         while True:
@@ -59,7 +67,7 @@ class RouterManager:
             self.recv_reqs.append(recv_req)
 
 
-def start_router_process(
+def start_controller_process(
     server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args
 ):
     logging.basicConfig(
@@ -68,8 +76,14 @@ def start_router_process(
     )
 
     try:
-        model_client = ModelRpcClient(server_args, port_args, model_overide_args)
-        router = RouterManager(model_client, port_args)
+        tp_size_local = server_args.tp_size // server_args.nnodes
+        model_client = ModelTpClient(
+            [i for _ in range(server_args.nnodes) for i in range(tp_size_local)],
+            server_args,
+            port_args.model_port_args[0],
+            model_overide_args,
+        )
+        controller = ControllerSingle(model_client, port_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
@@ -77,6 +91,12 @@ def start_router_process(
     pipe_writer.send("init ok")
 
     loop = asyncio.new_event_loop()
+    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
     asyncio.set_event_loop(loop)
-    loop.create_task(router.loop_for_recv_requests())
-    loop.run_until_complete(router.loop_for_forward())
+    loop.create_task(controller.loop_for_recv_requests())
+    try:
+        loop.run_until_complete(controller.loop_for_forward())
+    except Exception:
+        logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
+    finally:
+        kill_parent_process()

sglang/srt/managers/{router → controller}/model_runner.py

@@ -1,69 +1,40 @@
+"""ModelRunner runs the forward passes of the models."""
+
 import importlib
 import importlib.resources
-import inspect
 import logging
 import pkgutil
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import List
+from typing import List, Optional, Type
 
 import numpy as np
 import torch
-from vllm.distributed import initialize_model_parallel
-from vllm.model_executor.layers.quantization.awq import AWQConfig
-from vllm.model_executor.layers.quantization.gptq import GPTQConfig
-from vllm.model_executor.layers.quantization.marlin import MarlinConfig
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
-
-from sglang.srt.managers.router.infer_batch import Batch, ForwardMode
+import torch.nn as nn
+from vllm.config import DeviceConfig, LoadConfig
+from vllm.config import ModelConfig as VllmModelConfig
+from vllm.distributed import init_distributed_environment, initialize_model_parallel
+from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.models import ModelRegistry
+
+from sglang.srt.managers.controller.infer_batch import Batch, ForwardMode
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
-from sglang.srt.utils import is_multimodal_model, get_available_gpu_memory
-
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import (
+    get_available_gpu_memory,
+    is_multimodal_model,
+    monkey_patch_vllm_dummy_weight_loader,
+    monkey_patch_vllm_p2p_access_check,
+)
 
-QUANTIZATION_CONFIG_MAPPING = {
-    "awq": AWQConfig,
-    "gptq": GPTQConfig,
-    "marlin": MarlinConfig,
-}
-
-logger = logging.getLogger("model_runner")
+logger = logging.getLogger("srt.model_runner")
 
 # for server args in model endpoints
 global_server_args_dict = {}
 
 
-@lru_cache()
-def import_model_classes():
-    model_arch_name_to_cls = {}
-    package_name = "sglang.srt.models"
-    package = importlib.import_module(package_name)
-    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
-        if not ispkg:
-            module = importlib.import_module(name)
-            if hasattr(module, "EntryClass"):
-                model_arch_name_to_cls[module.EntryClass.__name__] = module.EntryClass
-    return model_arch_name_to_cls
-
-
-def get_model_cls_by_arch_name(model_arch_names):
-    model_arch_name_to_cls = import_model_classes()
-
-    model_class = None
-    for arch in model_arch_names:
-        if arch in model_arch_name_to_cls:
-            model_class = model_arch_name_to_cls[arch]
-            break
-    else:
-        raise ValueError(
-            f"Unsupported architectures: {arch}. "
-            f"Supported list: {list(model_arch_name_to_cls.keys())}"
-        )
-    return model_class
-
-
 @dataclass
 class InputMetadata:
-    model_runner: "ModelRunner"
     forward_mode: ForwardMode
     batch_size: int
     total_num_tokens: int
@@ -94,73 +65,82 @@ class InputMetadata:
     kv_indptr: torch.Tensor = None
     kv_indices: torch.Tensor = None
     kv_last_page_len: torch.Tensor = None
-    prefill_wrapper = None
-    decode_wrapper = None
+    flashinfer_prefill_wrapper_ragged: "BatchPrefillWithRaggedKVCacheWrapper" = None
+    flashinfer_prefill_wrapper_paged: "BatchPrefillWithPagedKVCacheWrapper" = None
+    flashinfer_decode_wrapper: "BatchDecodeWithPagedKVCacheWrapper" = None
 
-    def init_flashinfer_args(self, tp_size):
-        from flashinfer import (
-            BatchDecodeWithPagedKVCacheWrapper,
-            BatchPrefillWithPagedKVCacheWrapper,
-        )
+    def init_flashinfer_args(self, num_qo_heads, num_kv_heads, head_dim):
+        if (
+            self.forward_mode == ForwardMode.PREFILL
+            or self.forward_mode == ForwardMode.EXTEND
+        ):
+            paged_kernel_lens = self.prefix_lens
+            self.no_prefix = torch.all(self.prefix_lens == 0)
+        else:
+            paged_kernel_lens = self.seq_lens
 
         self.kv_indptr = torch.zeros(
             (self.batch_size + 1,), dtype=torch.int32, device="cuda"
         )
-        self.kv_indptr[1:] = torch.cumsum(self.seq_lens, dim=0)
+        self.kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
         self.kv_last_page_len = torch.ones(
             (self.batch_size,), dtype=torch.int32, device="cuda"
         )
         req_pool_indices_cpu = self.req_pool_indices.cpu().numpy()
-        seq_lens_cpu = self.seq_lens.cpu().numpy()
+        paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
         self.kv_indices = torch.cat(
             [
                 self.req_to_token_pool.req_to_token[
-                    req_pool_indices_cpu[i], : seq_lens_cpu[i]
+                    req_pool_indices_cpu[i], : paged_kernel_lens_cpu[i]
                 ]
                 for i in range(self.batch_size)
             ],
             dim=0,
         ).contiguous()
 
-        workspace_buffer = torch.empty(
-            32 * 1024 * 1024, dtype=torch.int8, device="cuda"
-        )
         if (
             self.forward_mode == ForwardMode.PREFILL
             or self.forward_mode == ForwardMode.EXTEND
         ):
+            # extend part
             self.qo_indptr = torch.zeros(
                 (self.batch_size + 1,), dtype=torch.int32, device="cuda"
             )
             self.qo_indptr[1:] = torch.cumsum(self.extend_seq_lens, dim=0)
-            self.prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
-                workspace_buffer, "NHD"
+
+            self.flashinfer_prefill_wrapper_ragged.end_forward()
+            self.flashinfer_prefill_wrapper_ragged.begin_forward(
+                self.qo_indptr,
+                self.qo_indptr.clone(),
+                num_qo_heads,
+                num_kv_heads,
+                head_dim,
             )
-            args = [
+
+            # cached part
+            self.flashinfer_prefill_wrapper_paged.end_forward()
+            self.flashinfer_prefill_wrapper_paged.begin_forward(
                 self.qo_indptr,
                 self.kv_indptr,
                 self.kv_indices,
                 self.kv_last_page_len,
-                self.model_runner.model_config.num_attention_heads // tp_size,
-                self.model_runner.model_config.num_key_value_heads // tp_size,
-                self.model_runner.model_config.head_dim,
-            ]
-
-            self.prefill_wrapper.begin_forward(*args)
-        else:
-            self.decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-                workspace_buffer, "NHD"
+                num_qo_heads,
+                num_kv_heads,
+                head_dim,
+                1
             )
-            self.decode_wrapper.begin_forward(
+        else:
+            self.flashinfer_decode_wrapper.end_forward()
+            self.flashinfer_decode_wrapper.begin_forward(
                 self.kv_indptr,
                 self.kv_indices,
                 self.kv_last_page_len,
-                self.model_runner.model_config.num_attention_heads // tp_size,
-                self.model_runner.model_config.num_key_value_heads // tp_size,
-                self.model_runner.model_config.head_dim,
+                num_qo_heads,
+                num_kv_heads,
+                head_dim,
                 1,
-                "NONE",
-                "float16",
+                pos_encoding_mode="NONE",
+                data_type=self.token_to_kv_pool.kv_data[0].dtype
             )
 
     def init_extend_args(self):
@@ -184,6 +164,9 @@ class InputMetadata:
         out_cache_cont_end=None,
         top_logprobs_nums=None,
         return_logprob=False,
+        flashinfer_prefill_wrapper_ragged=None,
+        flashinfer_prefill_wrapper_paged=None,
+        flashinfer_decode_wrapper=None,
     ):
         batch_size = len(req_pool_indices)
         start_loc = torch.zeros((batch_size,), dtype=torch.int32, device="cuda")
@@ -216,7 +199,6 @@ class InputMetadata:
             other_kv_index = None
 
         ret = cls(
-            model_runner=model_runner,
             forward_mode=forward_mode,
             batch_size=batch_size,
             total_num_tokens=total_num_tokens,
@@ -234,13 +216,20 @@ class InputMetadata:
             other_kv_index=other_kv_index,
             return_logprob=return_logprob,
             top_logprobs_nums=top_logprobs_nums,
+            flashinfer_prefill_wrapper_ragged=flashinfer_prefill_wrapper_ragged,
+            flashinfer_prefill_wrapper_paged=flashinfer_prefill_wrapper_paged,
+            flashinfer_decode_wrapper=flashinfer_decode_wrapper,
         )
 
         if forward_mode == ForwardMode.EXTEND:
             ret.init_extend_args()
 
-        if global_server_args_dict.get("enable_flashinfer", False):
-            ret.init_flashinfer_args(tp_size)
+        if not global_server_args_dict.get("disable_flashinfer", False):
+            ret.init_flashinfer_args(
+                model_runner.model_config.num_attention_heads // tp_size,
+                model_runner.model_config.get_num_kv_heads(tp_size),
+                model_runner.model_config.head_dim
+            )
 
         return ret
 
@@ -249,122 +238,180 @@ class ModelRunner:
     def __init__(
         self,
         model_config,
-        mem_fraction_static,
-        tp_rank,
-        tp_size,
-        nccl_port,
-        load_format="auto",
-        trust_remote_code=True,
-        server_args_dict: dict = {},
+        mem_fraction_static: float,
+        gpu_id: int,
+        tp_rank: int,
+        tp_size: int,
+        nccl_port: int,
+        server_args: ServerArgs,
     ):
         self.model_config = model_config
         self.mem_fraction_static = mem_fraction_static
+        self.gpu_id = gpu_id
         self.tp_rank = tp_rank
         self.tp_size = tp_size
         self.nccl_port = nccl_port
-        self.load_format = load_format
-        self.trust_remote_code = trust_remote_code
-
-        global global_server_args_dict
-        global_server_args_dict = server_args_dict
+        self.server_args = server_args
+        self.is_multimodal_model = is_multimodal_model(self.model_config)
+        monkey_patch_vllm_dummy_weight_loader()
 
         # Init torch distributed
-        torch.cuda.set_device(self.tp_rank)
-        torch.distributed.init_process_group(
+        logger.info(f"[gpu_id={self.gpu_id}] Set cuda device.")
+        torch.cuda.set_device(self.gpu_id)
+        logger.info(f"[gpu_id={self.gpu_id}] Init nccl begin.")
+        monkey_patch_vllm_p2p_access_check(self.gpu_id)
+        if server_args.nccl_init_addr:
+            nccl_init_method = f"tcp://{server_args.nccl_init_addr}"
+        else:
+            nccl_init_method = f"tcp://127.0.0.1:{self.nccl_port}"
+        init_distributed_environment(
             backend="nccl",
             world_size=self.tp_size,
             rank=self.tp_rank,
-            init_method=f"tcp://127.0.0.1:{self.nccl_port}",
+            local_rank=self.gpu_id,
+            distributed_init_method=nccl_init_method
         )
-
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
-
         total_gpu_memory = get_available_gpu_memory(
-            self.tp_rank, distributed=self.tp_size > 1
-        ) * (1 << 30)
+            self.gpu_id, distributed=self.tp_size > 1
+        )
+
+        if self.tp_size > 1:
+            total_local_gpu_memory = get_available_gpu_memory(self.gpu_id)
+            if total_local_gpu_memory < total_gpu_memory * 0.9:
+                raise ValueError(
+                    "The memory capacity is unbalanced. Some GPUs may be occupied by other processes."
+                )
+
+        # Set some global args
+        global global_server_args_dict
+        global_server_args_dict = {
+            "disable_flashinfer": server_args.disable_flashinfer,
+            "attention_reduce_in_fp32": server_args.attention_reduce_in_fp32,
+        }
+
+        # Load the model and create memory pool
         self.load_model()
         self.init_memory_pool(total_gpu_memory)
-
-        self.is_multimodal_model = is_multimodal_model(self.model_config)
+        self.init_cublas()
+        self.init_flash_infer()
 
     def load_model(self):
-        """See also vllm/model_executor/model_loader.py::get_model"""
-        # Select model class
-        architectures = getattr(self.model_config.hf_config, "architectures", [])
-        model_class = get_model_cls_by_arch_name(architectures)
-        logger.info(f"Rank {self.tp_rank}: load weight begin.")
-
-        # Load weights
-        quant_config = None
-
-        quant_cfg = getattr(self.model_config.hf_config, "quantization_config", None)
-        if quant_cfg is not None:
-            quant_method = quant_cfg.get("quant_method", "").lower()
-            # compat: autogptq >=0.8.0 use checkpoint_format: str
-            # compat: autogptq <=0.7.1 is_marlin_format: bool
-            is_format_marlin = quant_cfg.get(
-                "checkpoint_format"
-            ) == "marlin" or quant_cfg.get("is_marlin_format", False)
-
-            # Use marlin if the GPTQ model is serialized in marlin format.
-            if quant_method == "gptq" and is_format_marlin:
-                quant_method = "marlin"
-
-            quant_config_class = QUANTIZATION_CONFIG_MAPPING.get(quant_method)
-
-            if quant_config_class is None:
-                raise ValueError(f"Unsupported quantization method: {quant_method}")
-
-            quant_config = quant_config_class.from_config(quant_cfg)
-            logger.info(f"quant_config: {quant_config}")
-
-        with set_default_torch_dtype(torch.float16):
-            with torch.device("cuda"):
-                model = model_class(
-                    config=self.model_config.hf_config, quant_config=quant_config
-                )
-            model.load_weights(
-                self.model_config.path,
-                cache_dir=None,
-                load_format=self.load_format,
-                revision=None,
-            )
-        self.model = model.eval()
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Load weight begin. "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
 
-        logger.info(f"Rank {self.tp_rank}: load weight end.")
+        device_config = DeviceConfig()
+        load_config = LoadConfig(load_format=self.server_args.load_format)
+        vllm_model_config = VllmModelConfig(
+            model=self.server_args.model_path,
+            quantization=self.server_args.quantization,
+            tokenizer=None,
+            tokenizer_mode=None,
+            trust_remote_code=self.server_args.trust_remote_code,
+            dtype=self.server_args.dtype,
+            seed=42,
+            skip_tokenizer_init=True,
+        )
+        self.dtype = vllm_model_config.dtype
+        if self.model_config.model_overide_args is not None:
+            vllm_model_config.hf_config.update(self.model_config.model_overide_args)
+
+        self.model = get_model(
+            model_config=vllm_model_config,
+            device_config=device_config,
+            load_config=load_config,
+            lora_config=None,
+            vision_language_config=None,
+            parallel_config=None,
+            scheduler_config=None,
+            cache_config=None,
+        )
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Load weight end. "
+            f"type={type(self.model).__name__}, "
+            f"dtype={self.dtype}, "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
 
     def profile_max_num_token(self, total_gpu_memory):
         available_gpu_memory = get_available_gpu_memory(
-            self.tp_rank, distributed=self.tp_size > 1
-        ) * (1 << 30)
+            self.gpu_id, distributed=self.tp_size > 1
+        )
         head_dim = self.model_config.head_dim
-        head_num = self.model_config.num_key_value_heads // self.tp_size
-        cell_size = head_num * head_dim * self.model_config.num_hidden_layers * 2 * 2
+        head_num = self.model_config.get_num_kv_heads(self.tp_size)
+        cell_size = head_num * head_dim * self.model_config.num_hidden_layers * 2 * torch._utils._element_size(self.dtype)
         rest_memory = available_gpu_memory - total_gpu_memory * (
             1 - self.mem_fraction_static
         )
-        max_num_token = int(rest_memory // cell_size)
+        max_num_token = int(rest_memory * (1 << 30) // cell_size)
         return max_num_token
 
     def init_memory_pool(self, total_gpu_memory):
-        self.max_total_num_token = self.profile_max_num_token(total_gpu_memory)
+        self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
 
-        if self.max_total_num_token <= 0:
+        if self.max_total_num_tokens <= 0:
             raise RuntimeError(
-                "Not enought memory. " "Please try to increase --mem-fraction-static."
+                "Not enough memory. Please try to increase --mem-fraction-static."
             )
 
         self.req_to_token_pool = ReqToTokenPool(
-            int(self.max_total_num_token / self.model_config.context_len * 256),
+            int(self.max_total_num_tokens / self.model_config.context_len * 256),
             self.model_config.context_len + 8,
         )
         self.token_to_kv_pool = TokenToKVPool(
-            self.max_total_num_token,
-            dtype=torch.float16,
-            head_num=self.model_config.num_key_value_heads // self.tp_size,
+            self.max_total_num_tokens,
+            dtype=self.dtype,
+            head_num=self.model_config.get_num_kv_heads(self.tp_size),
             head_dim=self.model_config.head_dim,
             layer_num=self.model_config.num_hidden_layers,
         )
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Memory pool end. "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
+
+    def init_cublas(self):
+        """We need to run a small matmul to init cublas. Otherwise, it will raise some errors later."""
+        dtype = torch.float16
+        device = "cuda"
+        a = torch.ones((16, 16), dtype=dtype, device=device)
+        b = torch.ones((16, 16), dtype=dtype, device=device)
+        c = a @ b
+        return c
+
+    def init_flash_infer(self):
+        if not global_server_args_dict.get("disable_flashinfer", False):
+            from flashinfer import (
+                BatchPrefillWithRaggedKVCacheWrapper,
+                BatchPrefillWithPagedKVCacheWrapper,
+                BatchDecodeWithPagedKVCacheWrapper,
+            )
+            from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+
+            if not _grouped_size_compiled_for_decode_kernels(
+                self.model_config.num_attention_heads // self.tp_size,
+                self.model_config.get_num_kv_heads(self.tp_size)):
+                use_tensor_cores = True
+            else:
+                use_tensor_cores = False
+
+            workspace_buffers = torch.empty(
+                3, 96 * 1024 * 1024, dtype=torch.uint8, device="cuda"
+            )
+            self.flashinfer_prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
+                workspace_buffers[0], "NHD"
+            )
+            self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
+                workspace_buffers[1], "NHD"
+            )
+            self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
+                workspace_buffers[2], "NHD", use_tensor_cores=use_tensor_cores
+            )
+        else:
+            self.flashinfer_prefill_wrapper_ragged = self.flashinfer_prefill_wrapper_paged = None
+            self.flashinfer_decode_wrapper = None
 
     @torch.inference_mode()
     def forward_prefill(self, batch: Batch):
@@ -379,6 +426,9 @@ class ModelRunner:
             out_cache_loc=batch.out_cache_loc,
             top_logprobs_nums=batch.top_logprobs_nums,
             return_logprob=batch.return_logprob,
+            flashinfer_prefill_wrapper_ragged=self.flashinfer_prefill_wrapper_ragged,
+            flashinfer_prefill_wrapper_paged=self.flashinfer_prefill_wrapper_paged,
+            flashinfer_decode_wrapper=self.flashinfer_decode_wrapper,
         )
         return self.model.forward(
             batch.input_ids, input_metadata.positions, input_metadata
@@ -397,6 +447,9 @@ class ModelRunner:
             out_cache_loc=batch.out_cache_loc,
             top_logprobs_nums=batch.top_logprobs_nums,
             return_logprob=batch.return_logprob,
+            flashinfer_prefill_wrapper_ragged=self.flashinfer_prefill_wrapper_ragged,
+            flashinfer_prefill_wrapper_paged=self.flashinfer_prefill_wrapper_paged,
+            flashinfer_decode_wrapper=self.flashinfer_decode_wrapper,
         )
         return self.model.forward(
             batch.input_ids, input_metadata.positions, input_metadata
@@ -417,6 +470,9 @@ class ModelRunner:
             out_cache_cont_end=batch.out_cache_cont_end,
             top_logprobs_nums=batch.top_logprobs_nums,
             return_logprob=batch.return_logprob,
+            flashinfer_prefill_wrapper_ragged=self.flashinfer_prefill_wrapper_ragged,
+            flashinfer_prefill_wrapper_paged=self.flashinfer_prefill_wrapper_paged,
+            flashinfer_decode_wrapper=self.flashinfer_decode_wrapper,
         )
         return self.model.forward(
             batch.input_ids, input_metadata.positions, input_metadata
@@ -435,6 +491,9 @@ class ModelRunner:
             out_cache_loc=batch.out_cache_loc,
             top_logprobs_nums=batch.top_logprobs_nums,
             return_logprob=batch.return_logprob,
+            flashinfer_prefill_wrapper_ragged=self.flashinfer_prefill_wrapper_ragged,
+            flashinfer_prefill_wrapper_paged=self.flashinfer_prefill_wrapper_paged,
+            flashinfer_decode_wrapper=self.flashinfer_decode_wrapper,
         )
         return self.model.forward(
             batch.input_ids,
@@ -456,3 +515,48 @@ class ModelRunner:
             return self.forward_prefill(batch)
         else:
             raise ValueError(f"Invaid forward mode: {forward_mode}")
+
+
+@lru_cache()
+def import_model_classes():
+    model_arch_name_to_cls = {}
+    package_name = "sglang.srt.models"
+    package = importlib.import_module(package_name)
+    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
+        if not ispkg:
+            module = importlib.import_module(name)
+            if hasattr(module, "EntryClass"):
+                entry = module.EntryClass
+                if isinstance(
+                    entry, list
+                ):  # To support multiple model classes in one module
+                    for tmp in entry:
+                        model_arch_name_to_cls[tmp.__name__] = tmp
+                else:
+                    model_arch_name_to_cls[entry.__name__] = entry
+
+            # compat: some models such as chatglm has incorrect class set in config.json
+            # usage: [ tuple("From_Entry_Class_Name": EntryClass), ]
+            if hasattr(module, "EntryClassRemapping") and isinstance(
+                module.EntryClassRemapping, list
+            ):
+                for remap in module.EntryClassRemapping:
+                    if isinstance(remap, tuple) and len(remap) == 2:
+                        model_arch_name_to_cls[remap[0]] = remap[1]
+
+    return model_arch_name_to_cls
+
+
+def load_model_cls_srt(model_arch: str) -> Optional[Type[nn.Module]]:
+    model_arch_name_to_cls = import_model_classes()
+
+    if model_arch not in model_arch_name_to_cls:
+        raise ValueError(
+            f"Unsupported architectures: {model_arch}. "
+            f"Supported list: {list(model_arch_name_to_cls.keys())}"
+        )
+    return model_arch_name_to_cls[model_arch]
+
+
+# Monkey patch model loader
+setattr(ModelRegistry, "load_model_cls", load_model_cls_srt)