feat(backend): add sglang lora params of gpu (#272)

Co-authored-by: wasamtc <wasam@qq.com>

Author: wasamtc

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
   "dijkstar==2.6.0",
   "lattica==1.0.13",
   "orjson",
+  "transformers==4.55.2",
 ]
 
 [project.scripts]

src/parallax/server/executor.py

@@ -102,6 +102,15 @@ def __init__(
         # GPU/SGLang Specialized Configs
         attention_backend: Optional[str] = "flashinfer",
         moe_runner_backend: Optional[str] = "auto",
+        enable_lora: Optional[bool] = False,
+        max_lora_rank: Optional[int] = None,
+        lora_target_modules: Optional[List[str]] = None,
+        lora_paths: Optional[List[str]] = None,
+        max_loras_per_batch: Optional[int] = None,
+        max_loaded_loras: Optional[int] = None,
+        lora_eviction_policy: Optional[str] = "lru",
+        lora_backend: Optional[str] = "triton",
+        max_lora_chunk_size: Optional[int] = 128,
         # Tensor Parallel Configs
         tp_rank: Optional[int] = 0,
         tp_size: Optional[int] = 1,
@@ -155,6 +164,15 @@ def __init__(
                 "tp_size": tp_size,
                 "nccl_port": nccl_port,
                 "using_hfcache": use_hfcache,
+                "enable_lora": enable_lora,
+                "max_lora_rank": max_lora_rank,
+                "lora_target_modules": lora_target_modules,
+                "lora_paths": lora_paths,
+                "max_loras_per_batch": max_loras_per_batch,
+                "max_loaded_loras": max_loaded_loras,
+                "lora_eviction_policy": lora_eviction_policy,
+                "lora_backend": lora_backend,
+                "max_lora_chunk_size": max_lora_chunk_size,
             }
 
             self.model_runner, self.config, self.tokenizer = initialize_cuda_model_runner(
@@ -1615,5 +1633,14 @@ def create_executor_config(args: argparse.Namespace, gradient_server=None):
         "nccl_port": args.nccl_port,
         "gradient_server": gradient_server,
         "use_hfcache": args.use_hfcache,
+        "enable_lora": args.enable_lora,
+        "max_lora_rank": args.max_lora_rank,
+        "lora_target_modules": args.lora_target_modules,
+        "lora_paths": args.lora_paths,
+        "max_loras_per_batch": args.max_loras_per_batch,
+        "max_loaded_loras": args.max_loaded_loras,
+        "lora_eviction_policy": args.lora_eviction_policy,
+        "lora_backend": args.lora_backend,
+        "max_lora_chunk_size": args.max_lora_chunk_size,
     }
     return config

src/parallax/server/server_args.py

@@ -171,6 +171,69 @@ def parse_args() -> argparse.Namespace:
         help="Choose the GPU moe kernels",
     )
 
+    parser.add_argument(
+        "--enable-lora", action="store_true", help="Enable LoRA adapter support for SGLang backend"
+    )
+
+    parser.add_argument(
+        "--max-lora-rank",
+        type=int,
+        default=None,
+        help="The maximum rank of LoRA adapters. If not specified, it will be automatically inferred from the adapters provided in --lora-paths.",
+    )
+
+    parser.add_argument(
+        "--lora-target-modules",
+        nargs="*",
+        type=str,
+        default=None,
+        help="The union set of all target modules where LoRA should be applied. If not specified, it will be automatically inferred from the adapters provided in --lora-paths. If 'all' is specified, all supported modules will be targeted.",
+    )
+
+    parser.add_argument(
+        "--lora-paths",
+        nargs="*",
+        type=str,
+        default=None,
+        help="The list of LoRA adapters to load. Each adapter must be specified in one of the following formats: <PATH> | <NAME>=<PATH> | JSON with schema {'lora_name':str,'lora_path':str,'pinned':bool}.",
+    )
+
+    parser.add_argument(
+        "--max-loras-per-batch",
+        type=int,
+        default=8,
+        help="Maximum number of adapters for a running batch, include base-only request.",
+    )
+
+    parser.add_argument(
+        "--max-loaded-loras",
+        type=int,
+        default=None,
+        help="If specified, it limits the maximum number of LoRA adapters loaded in CPU memory at a time. The value must be greater than or equal to --max-loras-per-batch.",
+    )
+
+    parser.add_argument(
+        "--lora-eviction-policy",
+        choices=["lru", "fifo"],
+        default="lru",
+        help="LoRA adapter eviction policy when memory pool is full. 'lru': Least Recently Used (default, better cache efficiency). 'fifo': First-In-First-Out.",
+    )
+
+    parser.add_argument(
+        "--lora-backend",
+        choices=["triton", "csgmv"],
+        default="triton",
+        help="Choose the kernel backend for multi-LoRA serving.",
+    )
+
+    parser.add_argument(
+        "--max-lora-chunk-size",
+        choices=[16, 32, 64, 128],
+        type=int,
+        default=16,
+        help="Maximum chunk size for the ChunkedSGMV LoRA backend. Only used when --lora-backend is 'csgmv'. Choosing a larger value might improve performance.",
+    )
+
     # Tensor parallel configuration
     parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
 

src/parallax/sglang/model_runner.py

@@ -7,6 +7,7 @@
 import logging
 import os
 import random
+from typing import List, Optional
 
 import sglang
 import sglang.srt.distributed.parallel_state
@@ -207,6 +208,15 @@ def form_sgl_server_args(
     attention_backend: str = "flashinfer",
     kv_block_size: int = 64,
     moe_runner_backend="auto",
+    enable_lora: Optional[bool] = False,
+    max_lora_rank: Optional[int] = None,
+    lora_target_modules: Optional[List[str]] = None,
+    lora_paths: Optional[List[str]] = None,
+    max_loras_per_batch: Optional[int] = None,
+    max_loaded_loras: Optional[int] = None,
+    lora_eviction_policy: Optional[str] = "lru",
+    lora_backend: Optional[str] = "triton",
+    max_lora_chunk_size: Optional[int] = 128,
 ):
     """Creates a SGL ServerArgs object"""
     sgl_server_args = ServerArgs(
@@ -218,6 +228,15 @@ def form_sgl_server_args(
         moe_runner_backend=moe_runner_backend,
         tp_size=tp_size,
         trust_remote_code=True,
+        enable_lora=enable_lora,
+        max_lora_rank=max_lora_rank,
+        lora_target_modules=lora_target_modules,
+        lora_paths=lora_paths,
+        max_loras_per_batch=max_loras_per_batch,
+        max_loaded_loras=max_loaded_loras,
+        lora_eviction_policy=lora_eviction_policy,
+        lora_backend=lora_backend,
+        max_lora_chunk_size=max_lora_chunk_size,
     )
     return sgl_server_args
 
@@ -231,6 +250,15 @@ def initialize_sgl_model_runner(
     kv_block_size: int,
     moe_runner_backend: str,
     max_num_tokens_per_batch: int = 1024,
+    enable_lora: Optional[bool] = False,
+    max_lora_rank: Optional[int] = None,
+    lora_target_modules: Optional[List[str]] = None,
+    lora_paths: Optional[List[str]] = None,
+    max_loras_per_batch: Optional[int] = None,
+    max_loaded_loras: Optional[int] = None,
+    lora_eviction_policy: Optional[str] = "lru",
+    lora_backend: Optional[str] = "triton",
+    max_lora_chunk_size: Optional[int] = 128,
     **kwargs,
 ):
     """
@@ -285,6 +313,15 @@ def initialize_sgl_model_runner(
         attention_backend,
         kv_block_size,
         moe_runner_backend,
+        enable_lora,
+        max_lora_rank,
+        lora_target_modules,
+        lora_paths,
+        max_loras_per_batch,
+        max_loaded_loras,
+        lora_eviction_policy,
+        lora_backend,
+        max_lora_chunk_size,
     )
     initialize_moe_config(server_args)
     quant_method = None

tests/test_server_args.py

@@ -71,6 +71,41 @@ class TestCreateExecutorConfig:
 
     def test_create_config(self):
         """Test creating executor configuration."""
+        args = argparse.Namespace(
+            model_path="mlx-community/Qwen3-0.6B-bf16",
+            start_layer=0,
+            end_layer=10,
+            dtype="float16",
+            gpu_backend="sglang",
+            max_sequence_length=2048,
+            max_batch_size=8,
+            kv_block_size=64,
+            kv_cache_memory_fraction=0.8,
+            enable_prefix_cache=False,
+            max_num_tokens_per_batch=1024,
+            prefill_priority=0,
+            micro_batch_ratio=2,
+            scheduler_wait_ms=500,
+            send_to_peer_addr=None,
+            recv_from_peer_addr=None,
+            executor_input_ipc="ipc://test_input",
+            executor_output_ipc="ipc://test_output",
+            attention_backend="flashinfer",
+            moe_runner_backend="auto",
+            tp_rank=0,
+            tp_size=1,
+            nccl_port=4000,
+            use_hfcache=False,
+            enable_lora=False,
+            max_lora_rank=None,
+            lora_target_modules=None,
+            lora_paths=None,
+            max_loras_per_batch=1,
+            max_loaded_loras=8,
+            lora_eviction_policy="lru",
+            lora_backend="triton",
+            max_lora_chunk_size=128,
+        )
         args = argparse.Namespace(
             model_path="mlx-community/Qwen3-0.6B-bf16",
             start_layer=0,
@@ -92,6 +127,15 @@ def test_create_config(self):
             tp_size=1,
             nccl_port=4001,
             use_hfcache=False,
+            enable_lora=False,
+            max_lora_rank=None,
+            lora_target_modules=None,
+            lora_paths=None,
+            max_loras_per_batch=1,
+            max_loaded_loras=8,
+            lora_eviction_policy="lru",
+            lora_backend="triton",
+            max_lora_chunk_size=128,
         )
 
         config = create_executor_config(args)