sglang 0.4.2.post3__py3-none-any.whl → 0.4.3__py3-none-any.whl

sglang/srt/lora/triton_ops/gate_up_lora_b.py

@@ -0,0 +1,170 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _gate_up_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    output_dim,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    # Meta parameters
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling and adding
+    fuse_scaling_add,
+    scaling,
+):
+    # This kernel packs 2 sgemms (gate/up) into a single kernel.
+
+    # x: (s, 2 * K), s is the sum of sequence lengths, K equals to lora rank
+    # weights: (num_lora, 2 * output_dim, K)
+    # output: (s, 2 * output_dim)
+    # output_dim >> K
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # gate_up_id decides which of gate or up (0: gate, 1: up)
+    batch_id = tl.program_id(axis=2)
+    gate_up_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    w_index = tl.load(weight_indices + batch_id)
+    seg_start = tl.load(seg_indptr + batch_id)
+    n_start = gate_up_id * output_dim  # offset on output dim
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(output_dim, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+
+    # Create pointers for the first block of x and weights
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    x_ptrs = (x + seg_start * x_stride_0 + (gate_up_id * K) * x_stride_1) + (
+        s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+    )
+    w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iteate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len)
+            and (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K)
+            and (n_offset[None, :] < output_dim),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + (
+        s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset[:, None] < seg_len) and (n_offset[None, :] < output_dim)
+    if fuse_scaling_add:
+        partial_sum += tl.load(output_ptr, mask=output_mask)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def gate_up_lora_b_fwd(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor = None,
+    scaling: float = 1.0,
+) -> torch.Tensor:
+
+    # x: (s, 2 * r)
+    # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+    # output: (s, 2 * output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
+    # lora_output[:, output_dim:]
+    #      = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = gate_up_lora_b.shape[-1]
+    assert input_dim == 2 * r
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
+        2,  # this dimension decides current block computes on gate or up proj
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.empty((s, 2 * output_dim), device=x.device, dtype=x.dtype)
+        fuse_scaling_add = False
+    else:
+        output = base_output
+        fuse_scaling_add = True
+
+    _gate_up_lora_b_kernel[grid_b](
+        x,
+        gate_up_lora_b,
+        output,
+        r,
+        output_dim,
+        x.stride(0),
+        x.stride(1),
+        gate_up_lora_b.stride(0),
+        gate_up_lora_b.stride(1),
+        gate_up_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        fuse_scaling_add,
+        scaling,
+    )
+
+    return output

sglang/srt/lora/triton_ops/qkv_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 @triton.jit
@@ -108,7 +108,7 @@ def _qkv_lora_b_kernel(
 def qkv_lora_b_fwd(
     x: torch.Tensor,
     qkv_lora_b: torch.Tensor,
-    batch_info: LoraBatchInfo,
+    batch_info: LoRABatchInfo,
     output_offset: torch.Tensor,
     max_qkv_out_dim: int,
     base_output: torch.Tensor = None,
@@ -123,11 +123,11 @@ def qkv_lora_b_fwd(
     # output: (s, output_dim_q + 2 * output_dim_kv)
 
     # Compute lora_output with shape (s, output_dim) as follows:
-    # lora_output[:, :output_dim_q] = sgemm(lora_output_a[:, :r], )
+    # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
     # lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
-    #      = sgemm(lora_output_a[:, r: 2 * r], kv_lora_b[0])
+    #      = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
     # lora_output[:, output_dim_q + output_dim_kv: ]
-    #      = sgemm(lora_output_a[:, 2 * r: 3 * r], kv_lora_b[1])
+    #      = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
 
     # Get dims
     s = x.shape[0]

sglang/srt/lora/triton_ops/sgemm_lora_a.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 @triton.jit
@@ -91,7 +91,7 @@ def _sgemm_lora_a_kernel(
 
 
 def sgemm_lora_a_fwd(
-    x: torch.Tensor, weights: torch.Tensor, batch_info: LoraBatchInfo
+    x: torch.Tensor, weights: torch.Tensor, batch_info: LoRABatchInfo
 ) -> torch.Tensor:
     # x: (s, input_dim)
     # weights: (num_lora, r, input_dim)

sglang/srt/lora/triton_ops/sgemm_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 @triton.jit
@@ -98,7 +98,7 @@ def _sgemm_lora_b_kernel(
 def sgemm_lora_b_fwd(
     x: torch.Tensor,
     weights: torch.Tensor,
-    batch_info: LoraBatchInfo,
+    batch_info: LoRABatchInfo,
     base_output: torch.Tensor = None,
     scaling: float = 1.0,
 ) -> torch.Tensor:

sglang/srt/lora/utils.py

@@ -0,0 +1,141 @@
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Set, Tuple
+
+import torch
+
+from sglang.srt.hf_transformers_utils import AutoConfig
+
+
+@dataclass
+class LoRABatchInfo:
+    # Batch size
+    bs: int
+
+    # Lengths of each sequence in shape (bs,)
+    seg_lens: torch.Tensor
+
+    # Indice pointers of each sequence in shape (bs + 1, )
+    seg_indptr: torch.Tensor
+
+    # Maximum sequence length of current batch
+    max_len: int
+
+    # The index of lora adapter used by each sequence, in shape (bs,)
+    weight_indices: torch.Tensor
+
+
+class LoRAType(Enum):
+    LORA_A = 0
+    LORA_B = 1
+
+
+def get_layer_id(name: str) -> int:
+    """
+    Extract integer id of layer from its name in string.
+    """
+    match = re.search(r"layers\.(\d+)\.", name)
+    if match is None:
+        return None
+    return int(match.group(1))
+
+
+def get_customized_names_from_hf_names(
+    hf_module_names: Set[str], base_model: torch.nn.Module
+) -> Set[str]:
+    """
+    This function takes in a set of huggingface style module names:
+         e.g., {"k_proj", "q_proj", "v_proj", "o_proj"}
+    and outputs a set of module names of customized sglang layers:
+         e.g., {"qkv_proj", "o_proj"}
+    """
+    if hasattr(base_model, "get_module_name"):
+        return {base_model.get_module_name(name) for name in hf_module_names}
+    else:
+        """
+        Fallback solution of mapping from config module name to module name in model class.
+        Please check if it aligns with your base model.
+        Please implement the function in the model class if it is not.
+        You can reference this function in llama.py.
+        """
+        params_mapping = {
+            "q_proj": "qkv_proj",
+            "k_proj": "qkv_proj",
+            "v_proj": "qkv_proj",
+            "gate_proj": "gate_up_proj",
+            "up_proj": "gate_up_proj",
+        }
+        return {params_mapping.get(name, name) for name in hf_module_names}
+
+
+def get_hidden_dim(
+    module_name: str, config: AutoConfig, base_model: torch.nn.Module
+) -> Tuple[int]:
+    """
+    Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+    """
+
+    if hasattr(base_model, "get_hidden_dim"):
+        return base_model.get_hidden_dim(module_name)
+    else:
+        """
+        WARNING: get_hidden_dim() is not defined,
+        which is used to get the hidden dim for different lora modules
+        Use the default one, but please check if it is correct for your model.
+        Please implement the function in the model class if it is not.
+        You can reference this function in llama.py.
+        """
+        if module_name in ["q_proj", "o_proj", "qkv_proj"]:
+            return config.hidden_size, config.hidden_size
+        elif module_name in ["kv_proj"]:
+            return config.hidden_size, config.hidden_size // (
+                config.num_attention_heads // config.num_key_value_heads
+            )
+        elif module_name == "gate_up_proj":
+            return config.hidden_size, config.intermediate_size
+        elif module_name == "down_proj":
+            return config.intermediate_size, config.hidden_size
+        else:
+            raise NotImplementedError()
+
+
+def get_stacked_name(name: str) -> Tuple[str]:
+    """
+    Mapping a target module name to (stacked name for Lora A, stacked name for Lora B)
+    """
+    params_mapping = {
+        "q_proj": ("qkv_proj", "q_proj"),
+        "k_proj": ("qkv_proj", "kv_proj"),
+        "v_proj": ("qkv_proj", "kv_proj"),
+        "gate_proj": ("gate_up_proj", "gate_up_proj"),
+        "up_proj": ("gate_up_proj", "gate_up_proj"),
+    }
+    return params_mapping.get(name, (name, name))
+
+
+def get_stacked_multiply(module_name: str) -> int:
+    """
+    Mapping a lora module name to its magnification at output dimension
+    """
+    stacked_rank = {
+        "qkv_proj": 3,
+        "kv_proj": 2,
+        "gate_up_proj": 2,
+    }
+    return stacked_rank[module_name] if module_name in stacked_rank else 1
+
+
+def get_weight_name(
+    target_name: str, lora_weight_names: Set[Tuple[str]], lora_type: LoRAType
+) -> Optional[str]:
+    """
+    target_name is name of a given module,
+    lora_weight_names is a set of lora stacked name pairs (see get_stacked_name method above)
+    If there is a weight name in lora_weight_names that can match target_name, return this name
+    Else return None
+    """
+    idx = 0 if lora_type == LoRAType.LORA_A else 1
+    for weight_name_pair in lora_weight_names:
+        if weight_name_pair[idx] in target_name:
+            return weight_name_pair[idx]

sglang/srt/managers/detokenizer_manager.py

@@ -210,6 +210,7 @@ class DetokenizerManager:
                     input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
                     output_top_logprobs_val=recv_obj.output_top_logprobs_val,
                     output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
+                    output_hidden_states=recv_obj.output_hidden_states,
                 )
             )
 

sglang/srt/managers/io_struct.py

@@ -371,6 +371,8 @@ class BatchTokenIDOut:
     output_top_logprobs_val: List[List]
     output_top_logprobs_idx: List[List]
 
+    output_hidden_states: List[List[float]]
+
 
 @dataclass
 class BatchStrOut:
@@ -397,6 +399,8 @@ class BatchStrOut:
     output_top_logprobs_val: List[List]
     output_top_logprobs_idx: List[List]
 
+    output_hidden_states: List[List[float]]
+
 
 @dataclass
 class BatchEmbeddingOut:

sglang/srt/managers/schedule_batch.py

@@ -65,6 +65,7 @@ global_server_args_dict = {
     "enable_dp_attention": ServerArgs.enable_dp_attention,
     "enable_ep_moe": ServerArgs.enable_ep_moe,
     "device": ServerArgs.device,
+    "enable_flashinfer_mla": ServerArgs.enable_flashinfer_mla,
 }
 
 logger = logging.getLogger(__name__)
@@ -315,6 +316,7 @@ class Req:
             self.output_token_logprobs_val = self.output_token_logprobs_idx = (
                 self.output_top_logprobs_val
             ) = self.output_top_logprobs_idx = None
+        self.hidden_states = []
 
         # Logprobs (internal values)
         # The tokens is prefilled but need to be considered as decode tokens
@@ -604,6 +606,9 @@ class ScheduleBatch:
     # Enable custom logit processor
     enable_custom_logit_processor: bool = False
 
+    # Return hidden states
+    return_hidden_states: bool = False
+
     @classmethod
     def init_new(
         cls,
@@ -615,6 +620,7 @@ class ScheduleBatch:
         enable_overlap: bool,
         spec_algorithm: SpeculativeAlgorithm,
         enable_custom_logit_processor: bool,
+        return_hidden_states: bool = False,
     ):
         return cls(
             reqs=reqs,
@@ -629,6 +635,7 @@ class ScheduleBatch:
             device=req_to_token_pool.device,
             spec_algorithm=spec_algorithm,
             enable_custom_logit_processor=enable_custom_logit_processor,
+            return_hidden_states=return_hidden_states,
         )
 
     def batch_size(self):
@@ -1196,9 +1203,15 @@ class ScheduleBatch:
             spec_algorithm=self.spec_algorithm,
             spec_info=self.spec_info,
             capture_hidden_mode=(
-                getattr(self.spec_info, "capture_hidden_mode", CaptureHiddenMode.NULL)
-                if self.spec_info
-                else CaptureHiddenMode.NULL
+                CaptureHiddenMode.FULL
+                if self.return_hidden_states
+                else (
+                    getattr(
+                        self.spec_info, "capture_hidden_mode", CaptureHiddenMode.NULL
+                    )
+                    if self.spec_info
+                    else CaptureHiddenMode.NULL
+                )
             ),
         )
 

sglang/srt/managers/scheduler.py

@@ -997,6 +997,7 @@ class Scheduler:
             self.enable_overlap,
             self.spec_algorithm,
             self.server_args.enable_custom_logit_processor,
+            self.server_args.return_hidden_states,
         )
         new_batch.prepare_for_extend()
 
@@ -1156,6 +1157,8 @@ class Scheduler:
                         logits_output.input_token_logprobs.tolist()
                     )
 
+            hidden_state_offset = 0
+
             # Check finish conditions
             logprob_pt = 0
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
@@ -1182,6 +1185,21 @@ class Scheduler:
                             i, req, logprob_pt, next_token_ids, logits_output
                         )
 
+                    if (
+                        self.server_args.return_hidden_states
+                        and logits_output.hidden_states is not None
+                    ):
+                        req.hidden_states.append(
+                            logits_output.hidden_states[
+                                hidden_state_offset : (
+                                    hidden_state_offset := hidden_state_offset
+                                    + len(req.origin_input_ids)
+                                )
+                            ]
+                            .cpu()
+                            .clone()
+                        )
+
                     if req.grammar is not None:
                         req.grammar.accept_token(next_token_id)
                         req.grammar.finished = req.finished()
@@ -1275,6 +1293,12 @@ class Scheduler:
                         logits_output.next_token_top_logprobs_idx[i]
                     )
 
+            if (
+                self.server_args.return_hidden_states
+                and logits_output.hidden_states is not None
+            ):
+                req.hidden_states.append(logits_output.hidden_states[i].cpu().clone())
+
             if req.grammar is not None:
                 req.grammar.accept_token(next_token_id)
                 req.grammar.finished = req.finished()
@@ -1398,6 +1422,7 @@ class Scheduler:
             completion_tokens = []
             cached_tokens = []
             spec_verify_ct = []
+            hidden_states = []
 
             if return_logprob:
                 input_token_logprobs_val = []
@@ -1464,6 +1489,8 @@ class Scheduler:
                         output_top_logprobs_val.append(req.output_top_logprobs_val)
                         output_top_logprobs_idx.append(req.output_top_logprobs_idx)
 
+                    hidden_states.append(req.hidden_states)
+
             # Send to detokenizer
             if rids:
                 self.send_to_detokenizer.send_pyobj(
@@ -1490,6 +1517,7 @@ class Scheduler:
                         input_top_logprobs_idx,
                         output_top_logprobs_val,
                         output_top_logprobs_idx,
+                        hidden_states,
                     )
                 )
         else:  # embedding or reward model
@@ -1553,6 +1581,7 @@ class Scheduler:
             self.enable_overlap,
             self.spec_algorithm,
             self.server_args.enable_custom_logit_processor,
+            self.server_args.return_hidden_states,
         )
         idle_batch.prepare_for_idle()
         return idle_batch

sglang/srt/managers/tokenizer_manager.py

@@ -796,6 +796,12 @@ class TokenizerManager:
                     }
                 )
 
+            if (
+                hasattr(recv_obj, "output_hidden_states")
+                and len(recv_obj.output_hidden_states[i]) > 0
+            ):
+                meta_info["hidden_states"] = recv_obj.output_hidden_states[i]
+
             if isinstance(recv_obj, BatchStrOut):
                 out_dict = {
                     "text": recv_obj.output_strs[i],

sglang/srt/managers/tp_worker_overlap_thread.py

@@ -156,6 +156,10 @@ class TpModelWorkerClient:
                     logits_output.input_token_logprobs = (
                         logits_output.input_token_logprobs.to("cpu", non_blocking=True)
                     )
+            if logits_output.hidden_states is not None:
+                logits_output.hidden_states = logits_output.hidden_states.to(
+                    "cpu", non_blocking=True
+                )
             next_token_ids = next_token_ids.to("cpu", non_blocking=True)
             copy_done.record()
 

sglang/srt/model_executor/cuda_graph_runner.py

@@ -33,6 +33,9 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
 )
+from sglang.srt.utils import is_hip
+
+is_hip_ = is_hip()
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -129,6 +132,8 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
         if bs <= model_runner.req_to_token_pool.size
         and bs <= server_args.cuda_graph_max_bs
     ]
+    if is_hip_:
+        capture_bs += [i * 8 for i in range(21, 33)]
     compile_bs = (
         [bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
         if server_args.enable_torch_compile
@@ -237,6 +242,7 @@ class CudaGraphRunner:
                 "1. disable cuda graph by --disable-cuda-graph\n"
                 "2. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
                 "3. disable torch compile by not using --enable-torch-compile\n"
+                "4. set --cuda-graph-max-bs to a smaller value (e.g., 32)\n"
                 "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
             )
 
@@ -348,7 +354,13 @@ class CudaGraphRunner:
             spec_algorithm=self.model_runner.spec_algorithm,
             spec_info=spec_info,
             capture_hidden_mode=(
-                spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
+                CaptureHiddenMode.FULL
+                if self.model_runner.server_args.return_hidden_states
+                else (
+                    spec_info.capture_hidden_mode
+                    if spec_info
+                    else CaptureHiddenMode.NULL
+                )
             ),
         )
 
@@ -462,8 +474,11 @@ class CudaGraphRunner:
                     ),
                     positions=None,
                     retrive_index=None,
+                    retrive_next_token=None,
+                    retrive_next_sibling=None,
                     retrive_cum_len=None,
                     draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
+                    spec_steps=self.model_runner.server_args.speculative_num_steps,
                     capture_hidden_mode=CaptureHiddenMode.FULL,
                 )
 

sglang/srt/model_executor/model_runner.py

@@ -67,6 +67,7 @@ from sglang.srt.utils import (
     monkey_patch_p2p_access_check,
     monkey_patch_vllm_gguf_config,
     set_cpu_offload_max_bytes,
+    set_cuda_arch,
 )
 
 logger = logging.getLogger(__name__)
@@ -110,8 +111,14 @@ class ModelRunner:
         ):
             # TODO: add MLA optimization on CPU
             if self.server_args.device != "cpu":
-                logger.info("MLA optimization is turned on. Use triton backend.")
-                self.server_args.attention_backend = "triton"
+                if server_args.enable_flashinfer_mla:
+                    logger.info(
+                        "FlashInfer MLA optimization is turned on. Use flashinfer backend for DeepseekV3ForCausalLM."
+                    )
+                    self.server_args.attention_backend = "flashinfer"
+                else:
+                    logger.info("MLA optimization is turned on. Use triton backend.")
+                    self.server_args.attention_backend = "triton"
 
         if self.server_args.enable_double_sparsity:
             logger.info(
@@ -169,6 +176,7 @@ class ModelRunner:
                 "enable_dp_attention": server_args.enable_dp_attention,
                 "enable_ep_moe": server_args.enable_ep_moe,
                 "device": server_args.device,
+                "enable_flashinfer_mla": server_args.enable_flashinfer_mla,
             }
         )
 
@@ -292,6 +300,8 @@ class ModelRunner:
                 if torch.cuda.get_device_capability()[1] < 5:
                     raise RuntimeError("SGLang only supports sm75 and above.")
 
+        set_cuda_arch()
+
         # Prepare the model config
         self.load_config = LoadConfig(
             load_format=self.server_args.load_format,