[Feature] Add LoRA support for embedding layers (#14177)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Beichen-Ma <bm685@cornell.edu>

python/sglang/srt/layers/logits_processor.py

@@ -837,7 +837,10 @@ class LogitsProcessor(nn.Module):
             )
             dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)
 
-        if hasattr(lm_head, "weight"):
+        if hasattr(lm_head, "set_lora") and hasattr(lm_head, "apply_lora"):
+            # This is a LoRA-wrapped module, use its forward method
+            logits = lm_head(hidden_states)
+        elif hasattr(lm_head, "weight"):
             if self.use_fp32_lm_head:
                 logits = torch.matmul(
                     hidden_states.to(torch.float32), lm_head.weight.to(torch.float32).T

python/sglang/srt/lora/backend/base_backend.py

@@ -19,6 +19,52 @@ class BaseLoRABackend:
         self.max_loras_per_batch = max_loras_per_batch
         self.device = device
 
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run LoRA A embedding lookup with CUDA graph support.
+
+        Args:
+            input_ids: token IDs with shape (s,), where s is the sum of all sequence lengths
+            weights: LoRA A embedding weights with shape (num_loras, rank, vocab_size)
+            vocab_size: base vocabulary size (tokens >= vocab_size are extra tokens)
+            extra_embeddings: extra token embeddings with shape (num_loras, num_extra_tokens, rank)
+            Only needed if there are added tokens beyond base vocabulary.
+
+        Returns:
+            result with shape (s, rank)
+        """
+        pass
+
+    def run_extra_token_embedding(
+        self,
+        input_ids: torch.Tensor,
+        output: torch.Tensor,
+        extra_embeddings: torch.Tensor,
+        vocab_size: int,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Apply extra token embeddings to output in-place.
+
+        Args:
+            input_ids: (s,) token IDs
+            output: (s, embed_dim) output tensor to be modified
+            extra_embeddings: (num_loras, num_extra_tokens, embed_dim) extra embeddings
+            vocab_size: base vocabulary size
+
+        Returns:
+            output: modified output tensor
+        """
+        raise NotImplementedError
+
     def run_lora_a_sgemm(
         self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
     ) -> torch.Tensor:

python/sglang/srt/lora/backend/triton_backend.py

@@ -2,6 +2,7 @@ import torch
 
 from sglang.srt.lora.backend.base_backend import BaseLoRABackend
 from sglang.srt.lora.triton_ops import (
+    embedding_lora_a_fwd,
     gate_up_lora_b_fwd,
     qkv_lora_b_fwd,
     sgemm_lora_a_fwd,
@@ -22,6 +23,24 @@ class TritonLoRABackend(BaseLoRABackend):
     ):
         super().__init__(max_loras_per_batch, device)
 
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run LoRA A embedding lookup using Triton kernel."""
+        return embedding_lora_a_fwd(
+            input_ids=input_ids,
+            weights=weights,
+            batch_info=self.batch_info,
+            vocab_size=vocab_size,
+            extra_embeddings=extra_embeddings,
+        )
+
     def run_lora_a_sgemm(
         self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
     ) -> torch.Tensor:
@@ -107,7 +126,7 @@ class TritonLoRABackend(BaseLoRABackend):
                 seg_lens=torch.full(
                     (max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
                 ),
-                seg_indptr=torch.empty(max_bs_in_cuda_graph + 1, dtype=torch.int32),
+                seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
                 max_len=num_tokens_per_bs,
                 weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
                 lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
@@ -161,7 +180,7 @@ class TritonLoRABackend(BaseLoRABackend):
             seg_lens = (
                 forward_batch.extend_seq_lens
                 if forward_batch.forward_mode.is_extend()
-                else torch.ones(bs, device=self.device)
+                else torch.ones(bs, dtype=torch.int32, device=self.device)
             )
             seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
             seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)

python/sglang/srt/lora/layers.py

@@ -1,4 +1,7 @@
+from typing import Optional
+
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 from sglang.srt.distributed import (
@@ -13,8 +16,12 @@ from sglang.srt.layers.linear import (
     QKVParallelLinear,
     RowParallelLinear,
 )
-from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 class BaseLayerWithLoRA(nn.Module):
@@ -45,11 +52,10 @@ class BaseLayerWithLoRA(nn.Module):
 
 class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
     """
-    Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation).
+    Vocab parallel embedding layer with LoRA support (simplified for TP=1, no extra tokens).
 
-    Note: The current version does not yet implement the LoRA functionality.
-    This class behaves exactly the same as the base VocabParallelEmbedding.
-    Future versions will integrate LoRA functionality to support efficient parameter fine-tuning.
+    For embedding layers: output = base_embedding(x) + lora_B @ lora_A[x]
+    where lora_A[x] is direct embedding lookup from lora_A weights.
     """
 
     def __init__(
@@ -59,6 +65,237 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
     ) -> None:
         super().__init__(base_layer, lora_backend)
         self.weight = base_layer.weight
+        self.embed_dim = base_layer.embedding_dim
+        self.vocab_size = base_layer.org_vocab_size
+
+        self.output_offset = torch.tensor(
+            [0, self.embed_dim],
+            dtype=torch.int32,
+            device=next(base_layer.parameters()).device,
+        )
+
+    def set_lora_info(
+        self,
+        new_embeddings_buffer: Optional[torch.Tensor],  # For extra tokens
+        embedding_A_buffer: torch.Tensor,
+        embedding_B_buffer: torch.Tensor,
+    ):
+        """Set LoRA buffers for embedding layer."""
+        self.set_lora = True
+        self.new_embeddings_buffer = new_embeddings_buffer
+        self.embedding_A_buffer = embedding_A_buffer  # (num_loras, rank, vocab_size)
+        self.embedding_B_buffer = embedding_B_buffer  # (num_loras, embed_dim, rank)
+
+    def apply_lora(
+        self, base_output: torch.Tensor, input_: torch.Tensor, batch_info
+    ) -> torch.Tensor:
+        """
+        Apply LoRA to base embedding output.
+        Formula: output = base_output + lora_B @ lora_A_embedding(input_)
+        """
+
+        # Efficient embedding lookup for LoRA A (already support extra token embedding process)
+        lora_a_output = self.run_lora_a_embedding(input_, batch_info)
+
+        # Apply LoRA B weights using backend
+        lora_output = self.lora_backend.run_lora_b_sgemm(
+            x=lora_a_output,
+            weights=self.embedding_B_buffer,
+            output_offset=self.output_offset,
+            base_output=base_output,
+        )
+        return lora_output
+
+    def run_lora_a_embedding(
+        self, input_: torch.Tensor, batch_info: LoRABatchInfo
+    ) -> torch.Tensor:
+        """
+        Apply LoRA A weights using efficient embedding lookup with CUDA graph support.
+        Maps tokens to their corresponding LoRA adapters internally.
+        It also includes added/extra token processing.
+        """
+        # Efficient embedding lookup for LoRA A (already support extra token embedding process)
+        lora_a_output = self.lora_backend.run_lora_a_embedding(
+            input_ids=input_,
+            weights=self.embedding_A_buffer,
+            vocab_size=self.vocab_size,
+            extra_embeddings=(
+                self.new_embeddings_buffer
+                if hasattr(self, "new_embeddings_buffer")
+                and self.new_embeddings_buffer is not None
+                else None
+            ),
+        )
+
+        return lora_a_output
+
+    def extra_token_embedding(
+        self, input_: torch.Tensor, base_output: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Need to impl:
+
+        Process extra tokens (tokens >= vocab_size) by looking up their embeddings
+        from the new_embeddings_buffer and replacing them in base_output.
+
+        Args:
+            input_: (s,) token IDs
+            base_output: (s, embed_dim) base embedding output to be modified in-place
+
+        Returns:
+            base_output: (s, embed_dim) modified input base_output (tensor[0,0,0,...]) with extra token embeddings
+        """
+        # return base_output
+        raise NotImplementedError(
+            "Error in sglang/python/sglang/srt/lora/layers.py - VocabParallelEmbeddingWithLoRA \n"
+            "Current SGLang codebase did not support tuned lora with extra/added tokens. \n"
+            "[TODO]: \n"
+            "1. Refer to this commit: https://github.com/yushengsu-thu/sglang/commit/90415211eee8a28a316de262583d4d33fa615d10#diff-191177438bcc223837963de63c005850371f8c8a860acb153b26744b66ecc623 to complete \n"
+            "2. And then you need to modified the en/decoder tokenizer - tokenizer_manager.py to support extra_token_embedding in-place. \n"
+        )
+
+    def forward(self, input_: torch.Tensor):
+        """
+        Forward pass with LoRA support and CUDA graph compatibility.
+
+        Extra tokens (tokens >= vocab_size) are now handled efficiently
+        in the backend's run_lora_a_embedding method.
+        """
+        batch_info = self.lora_backend.batch_info
+
+        # Get base embedding output
+        # For tokens >= vocab_size, base_layer will clamp or handle them
+        # We mask them to 0 to avoid out-of-bounds access
+        added_tokens_mask = input_ > self.vocab_size - 1
+        base_output = self.base_layer.forward(input_.masked_fill(added_tokens_mask, 0))
+
+        # [TODO] SGLang did not support extra/added token process; thus, self.extra_token_embedding only return original input_ now
+        # Extra tokens - It will replace extra token embedding with self.new_embeddings_buffer's emb (Default is 0)
+        if (
+            hasattr(self, "new_embeddings_buffer")
+            and self.new_embeddings_buffer is not None
+        ):
+            base_output = self.extra_token_embedding(input_, base_output)
+
+        # Apply LoRA if configured
+        if self.set_lora:
+            # The backend's run_lora_a_embedding now handles both regular
+            # and extra tokens efficiently with CUDA graph support
+            base_output = self.apply_lora(base_output, input_, batch_info)
+
+        return base_output
+
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        # For TP=1, no slicing needed
+        # LoRA A weights (rank, vocab_size) are not sliced for embedding
+        # For TP>1, Need to modify code in: sglang/python/sglang/srt/lora/mem_pool.py
+        # return A
+        if tp_rank > 1:
+            raise NotImplementedError(
+                f"VocabParallelEmbeddingWithLoRA does not support tensor parallelism > 1. "
+                f"Got tp_size={tp_rank}"
+            )
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # For TP=1, no slicing needed
+        # LoRA B weights (embedding_dim, rank) would be sliced along embedding dimension for TP>1
+        # For TP>1, Need to modify code in: sglang/python/sglang/srt/lora/mem_pool.py
+        # return B
+        if tp_rank > 1:
+            raise NotImplementedError(
+                f"VocabParallelEmbeddingWithLoRA does not support tensor parallelism > 1. "
+                f"Got tp_size={tp_rank}"
+            )
+
+
+class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
+    """
+    Parallel LM Head layer with LoRA support (simplified for TP=1).
+
+    The LM head computes logits = hidden_states @ (W + B @ A)^T
+    """
+
+    def __init__(
+        self,
+        base_layer: ParallelLMHead,
+        lora_backend: BaseLoRABackend,
+    ) -> None:
+        super().__init__(base_layer, lora_backend)
+        self.weight = base_layer.weight
+        self.embed_dim = base_layer.embedding_dim
+        self.vocab_size = base_layer.org_vocab_size
+        self.output_offset = torch.tensor(
+            [0, self.vocab_size],
+            dtype=torch.int32,
+            device=next(base_layer.parameters()).device,
+        )
+
+    def set_lora_info(
+        self,
+        lm_head_A_buffer: torch.Tensor,
+        lm_head_B_buffer: torch.Tensor,
+    ):
+        """Set LoRA buffers for LM head layer."""
+        self.set_lora = True
+        self.lm_head_A_buffer = lm_head_A_buffer  # (num_loras, rank, hidden_dim)
+        self.lm_head_B_buffer = lm_head_B_buffer  # (num_loras, vocab_size, rank)
+
+    def apply_lora(
+        self, base_output: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Apply LoRA to LM head layer.
+
+        For LM head: output = hidden @ (W + B @ A)^T
+                           = hidden @ W^T + hidden @ A^T @ B^T
+                           = base_output + (hidden @ A^T) @ B^T
+        """
+        # Apply lora_A^T: hidden_states @ A^T
+        lora_a_output = self.lora_backend.run_lora_a_sgemm(
+            hidden_states, self.lm_head_A_buffer
+        )
+
+        # Apply lora_B^T: lora_a_output @ B^T
+        lora_output = self.lora_backend.run_lora_b_sgemm(
+            x=lora_a_output,
+            weights=self.lm_head_B_buffer,
+            output_offset=self.output_offset,
+            base_output=base_output,
+        )
+
+        return lora_output
+
+    def forward(self, hidden_states: torch.Tensor):
+        # Apply base linear transformation
+        base_output = F.linear(
+            hidden_states, self.weight, bias=getattr(self.base_layer, "bias", None)
+        )
+
+        # Apply LoRA if set
+        if self.set_lora:
+            base_output = self.apply_lora(base_output, hidden_states)
+
+        return base_output
+
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        # For TP=1, no slicing needed
+        # For TP>1, need to modify code in: sglang/python/sglang/srt/lora/mem_pool.py
+        # return A
+        if tp_rank > 1:
+            raise NotImplementedError(
+                f"ParallelLMHeadWithLoRA does not support tensor parallelism > 1. "
+                f"Got tp_size={tp_rank}"
+            )
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # For TP=1, no slicing needed
+        # For TP>1, would slice along vocab dimension, need to modify code in: sglang/python/sglang/srt/lora/mem_pool.py
+        # return B
+        if tp_rank > 1:
+            raise NotImplementedError(
+                f"ParallelLMHeadWithLoRA does not support tensor parallelism > 1. "
+                f"Got tp_size={tp_rank}"
+            )
 
 
 class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
@@ -224,6 +461,7 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
             output_offset_cpu=self.output_offset_cpu,
             max_qkv_out_dim=self.max_qkv_out_dim,
         )
+
         return lora_output
 
     def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
@@ -343,6 +581,7 @@ def get_lora_layer(
 ) -> BaseLayerWithLoRA:
     supported_layer_types = {
         # the order matters
+        ParallelLMHead: ParallelLMHeadWithLoRA,
         VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
         QKVParallelLinear: QKVParallelLinearWithLoRA,
         MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA,

python/sglang/srt/lora/lora.py

@@ -71,11 +71,22 @@ class LoRAAdapter(nn.Module):
             ]
         )
 
+        self.embedding_layers: Dict[str, torch.Tensor] = {}
+        self.added_tokens_embeddings: Dict[str, torch.Tensor] = {}
+
     # initialize the LoRA weights to cpu
     def initialize_weights(self):
         model_path = self.config.path
         loader = DefaultModelLoader(self.load_config)
         revision = getattr(self.config.hf_config, "revision", None)
+
+        # Get normalized target modules for filtering
+        from sglang.srt.lora.utils import get_normalized_target_modules
+
+        normalized_target_modules = get_normalized_target_modules(
+            self.config.target_modules
+        )
+
         for name, loaded_weight in loader._get_weights_iterator(
             DefaultModelLoader.Source(
                 model_path, revision=revision, fall_back_to_pt=True
@@ -84,6 +95,22 @@ class LoRAAdapter(nn.Module):
             layer_id = get_layer_id(name)
             if layer_id is not None:
                 self.layers[layer_id].weights[name] = loaded_weight.cpu()
+            elif "embed_tokens" in name or "lm_head" in name:
+                # Check if this module is declared in target_modules before loading
+                module_name = "embed_tokens" if "embed_tokens" in name else "lm_head"
+                if module_name in normalized_target_modules:
+                    self.embedding_layers[name] = loaded_weight.cpu()
+                else:
+                    logger.debug(
+                        f"Skipping {name} as '{module_name}' is not in adapter's target_modules: {self.config.target_modules}"
+                    )
+            elif "input_embeddings" in name or "output_embeddings" in name:
+                # added/extra token emb
+                self.added_tokens_embeddings[name] = loaded_weight.cpu()
+                assert loaded_weight.shape[0] == self.config.lora_added_tokens_size, (
+                    f"LoRA adapter {self.uid} has extra_vocab_size {self.config.extra_vocab_size} specified in the config, "
+                    f"but the loaded weight has {loaded_weight.shape[0]} extra vocab size"
+                )
 
         # normalize kv_proj and gate_up_proj
         for layer in self.layers:

python/sglang/srt/lora/lora_config.py

@@ -27,13 +27,14 @@ class LoRAConfig:
         self.hf_config = self.get_lora_config()
         self.target_modules = self.hf_config["target_modules"]
 
-        # TODO: Support more modules
-        if any(module in self.target_modules for module in ["embed_tokens", "lm_head"]):
-            raise ValueError("Not supported yet")
-
         self.r = self.hf_config["r"]
         self.lora_alpha = self.hf_config["lora_alpha"]
 
+        self.added_tokens_config = self.get_added_tokens_config()
+        self.lora_added_tokens_size = (
+            len(self.added_tokens_config) if self.added_tokens_config is not None else 0
+        )
+
     def get_lora_config(self, dummy=False):
         if dummy:
             raise NotImplementedError()
@@ -45,3 +46,30 @@ class LoRAConfig:
             config_name = "adapter_config.json"
             with open(os.path.join(weights_dir, config_name), "r") as f:
                 return json.load(f)
+
+    def get_added_tokens_config(self):
+        """Load added tokens from the LoRA adapter if the file exists."""
+        # Determine the weights directory
+        if not os.path.isdir(self.path):
+            weights_dir = snapshot_download(self.path, allow_patterns=["*.json"])
+        else:
+            weights_dir = self.path
+
+        # Construct the path to added_tokens.json
+        added_tokens_path = os.path.join(weights_dir, "added_tokens.json")
+
+        # Return None if the file doesn't exist (optional for standard LoRA adapters)
+        if not os.path.exists(added_tokens_path):
+            return None
+
+        # Load and return the added tokens
+        try:
+            with open(added_tokens_path, "r") as f:
+                return json.load(f)
+        except json.JSONDecodeError as e:
+            # Log warning but don't crash if JSON is malformed
+            import logging
+
+            logger = logging.getLogger(__name__)
+            logger.warning(f"Failed to parse added_tokens.json: {e}")
+            return None

python/sglang/srt/lora/lora_manager.py

@@ -22,6 +22,10 @@ import torch
 
 from sglang.srt.configs.load_config import LoadConfig
 from sglang.srt.layers.utils import get_layer_id
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from sglang.srt.lora.backend.base_backend import BaseLoRABackend
 from sglang.srt.lora.backend.lora_registry import get_backend_from_name
 from sglang.srt.lora.layers import BaseLayerWithLoRA, get_lora_layer
@@ -67,6 +71,7 @@ class LoRAManager:
         self.device: torch.device = next(self.base_model.parameters()).device
         self.tp_size: int = tp_size
         self.tp_rank: int = tp_rank
+        self.lora_added_tokens_size: Optional[int] = None
 
         # Store eviction policy from server args
         self.eviction_policy = server_args.lora_eviction_policy
@@ -244,6 +249,8 @@ class LoRAManager:
             lora_adapters=self.loras,
             lora_modules=self.lora_modules,
             lora_refs=self.lora_refs.copy(),  # copy snapshot of current lora_refs to avoid mutation during the batch preparation.
+            lora_embed_tokens_module=self.embed_tokens_module,  # merge into embedding or lora module
+            lora_lm_head_module=self.lm_head_module,  # merge into embedding or lora module
         )
 
         # set up batch info shared by all lora modules
@@ -296,6 +303,21 @@ class LoRAManager:
                     ),
                 )
 
+        # Update embedding layer if present - gotta merge (refer to PR codebase)
+        if self.embed_tokens_module is not None:
+            self.embed_tokens_module.set_lora_info(
+                self.memory_pool.get_embedding_tensor("added_tokens", LoRAType.LORA_A),
+                self.memory_pool.get_embedding_tensor("embed_tokens", LoRAType.LORA_A),
+                self.memory_pool.get_embedding_tensor("embed_tokens", LoRAType.LORA_B),
+            )
+
+        # Update lm_head layer if present
+        if self.lm_head_module is not None:
+            self.lm_head_module.set_lora_info(
+                self.memory_pool.get_embedding_tensor("lm_head", LoRAType.LORA_A),
+                self.memory_pool.get_embedding_tensor("lm_head", LoRAType.LORA_B),
+            )
+
     def init_state(
         self,
         max_lora_rank: Optional[int] = None,
@@ -390,6 +412,24 @@ class LoRAManager:
                 default=0,
             )
 
+        # Auto-infer self.lora_added_vocab_size from loaded LoRA configs
+        # This happens automatically without requiring user input
+        # if self.lora_added_vocab_size is None:
+        if self.lora_added_tokens_size is None:
+            inferred_extra_vocab_size = next(
+                (
+                    x.lora_added_tokens_size
+                    for x in self.configs.values()
+                    if x.lora_added_tokens_size > 0
+                ),
+                0,
+            )
+            if inferred_extra_vocab_size > 0:
+                logger.info(
+                    f"self.lora_added_tokens_size={inferred_extra_vocab_size} from LoRA adapters."
+                )
+            self.lora_added_tokens_size = inferred_extra_vocab_size
+
     def load_lora_weights(self, lora_ref: LoRARef):
         """
         Load the weights of a LoRA adapter to CPU memory and conducts post-loading validation.
@@ -416,6 +456,7 @@ class LoRAManager:
             target_modules=self.target_modules,
             base_model=self.base_model,
             eviction_policy=self.eviction_policy,
+            lora_added_tokens_size=self.lora_added_tokens_size,
         )
 
     def set_lora_module(self, module_name, module):
@@ -429,6 +470,9 @@ class LoRAManager:
             {} for _ in range(self.base_hf_config.num_hidden_layers)
         ]
 
+        self.embed_tokens_module: Optional[BaseLayerWithLoRA] = None
+        self.lm_head_module: Optional[BaseLayerWithLoRA] = None
+
         for module_name, module in self.base_model.named_modules():
             # TODO (lifuhuang): in the future, we should consider generalizing the
             # should_apply_lora function to support mapping by full module name instead
@@ -440,6 +484,24 @@ class LoRAManager:
             ) and not self.base_model.should_apply_lora(module_name):
                 continue
 
+            # Handle embed_tokens
+            if "embed_tokens" in module_name and "embed_tokens" in self.target_modules:
+                if isinstance(module, VocabParallelEmbedding) and not isinstance(
+                    module, BaseLayerWithLoRA
+                ):
+                    lora_module = self.set_lora_module(module_name, module)
+                    self.embed_tokens_module = lora_module
+                    continue
+
+            # Handle lm_head
+            if "lm_head" in module_name and "lm_head" in self.target_modules:
+                if isinstance(module, ParallelLMHead) and not isinstance(
+                    module, BaseLayerWithLoRA
+                ):
+                    lora_module = self.set_lora_module(module_name, module)
+                    self.lm_head_module = lora_module
+                    continue
+
             # The module should be converted if it is included in target_names
             if module_name.split(".")[-1] in self.target_modules:
                 layer_id = get_layer_id(module_name)

python/sglang/srt/lora/mem_pool.py

@@ -10,6 +10,7 @@ from sglang.srt.lora.lora import LoRAAdapter
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.lora.utils import (
+    EMBEDDING_NAMES,
     ROW_PARALLELISM_LINEAR_LORA_NAMES,
     LoRAType,
     get_hidden_dim,
@@ -56,6 +57,7 @@ class LoRAMemoryPool:
         target_modules: Set[str],
         base_model: torch.nn.Module,
         eviction_policy: str,
+        lora_added_tokens_size: int,
     ):
         self.base_hf_config: AutoConfig = base_hf_config
         self.num_layer: int = base_hf_config.num_hidden_layers
@@ -63,6 +65,7 @@ class LoRAMemoryPool:
         self.dtype: torch.dtype = dtype
         self.tp_size: int = tp_size
         self.tp_rank: int = tp_rank
+        self.lora_added_tokens_size: int = lora_added_tokens_size
         self.max_lora_rank: int = max_lora_rank
         self.target_modules: Set[str] = target_modules
 
@@ -77,6 +80,15 @@ class LoRAMemoryPool:
         self.A_buffer: Dict[str, List[torch.Tensor]] = {}
         self.B_buffer: Dict[str, List[torch.Tensor]] = {}
 
+        self.embedding_A_buffer: Dict[str, torch.Tensor] = {}
+        self.embedding_B_buffer: Dict[str, torch.Tensor] = {}
+
+        self.lm_head_A_buffer: Dict[str, torch.Tensor] = {}
+        self.lm_head_B_buffer: Dict[str, torch.Tensor] = {}
+        self.new_embeddings_buffer: Dict[str, torch.Tensor] = {}
+
+        self.embedding_dim: int = self.base_hf_config.hidden_size
+
         # Lora uid -> buffer idx in memory pool
         self.uid_to_buffer_id: Dict[Optional[str], int] = {}
 
@@ -100,6 +112,8 @@ class LoRAMemoryPool:
             """
             if config.r > self.max_lora_rank:
                 return False
+            if config.lora_added_tokens_size > self.lora_added_tokens_size:
+                return False
             target_module_names = get_normalized_target_modules(config.target_modules)
             return target_module_names.issubset(self.target_modules)
 
@@ -130,6 +144,23 @@ class LoRAMemoryPool:
             input_dim,
         )
 
+    def get_embedding_lora_A_shape(
+        self,
+        module_name: str,
+        base_model: torch.nn.Module,
+        max_lora_dim: int,
+        layer_idx: int,
+    ) -> Tuple[int]:
+        input_dim, _ = get_hidden_dim(
+            module_name, self.base_hf_config, base_model, 0, self.lora_added_tokens_size
+        )
+        # Have not imp self.tp_size > 1 yet.
+        return (
+            self.max_loras_per_batch,
+            max_lora_dim,
+            input_dim,
+        )
+
     def get_lora_B_shape(
         self,
         module_name: str,
@@ -151,6 +182,23 @@ class LoRAMemoryPool:
             max_lora_dim,
         )
 
+    def get_embedding_lora_B_shape(
+        self,
+        module_name: str,
+        base_model: torch.nn.Module,
+        max_lora_dim: int,
+        layer_idx: int,
+    ) -> Tuple[int]:
+        _, output_dim = get_hidden_dim(
+            module_name, self.base_hf_config, base_model, 0, self.lora_added_tokens_size
+        )
+        # Have not imp self.tp_size > 1 yet.
+        return (
+            self.max_loras_per_batch,
+            output_dim,
+            max_lora_dim,
+        )
+
     def init_buffers(self, base_model: torch.nn.Module):
         device = next(base_model.parameters()).device
 
@@ -159,6 +207,7 @@ class LoRAMemoryPool:
             target_modules: Set[str],
             get_lora_shape_fn: Callable[[str, torch.nn.Module, int, int], Tuple[int]],
         ):
+            target_modules = target_modules - set(EMBEDDING_NAMES)
             for module_name in target_modules:
                 buffer[module_name] = [
                     torch.empty(
@@ -174,6 +223,61 @@ class LoRAMemoryPool:
                     for idx in range(self.num_layer)
                 ]
 
+        def init_embedding_buffer(
+            buffer: Dict[str, torch.Tensor],
+            target_modules: Set[str],
+            get_lora_shape_fn: Callable[[int], Tuple[int]],
+        ):
+            target_modules = target_modules & set(EMBEDDING_NAMES)
+            for module_name in target_modules:
+                buffer[module_name] = torch.empty(
+                    get_lora_shape_fn(
+                        module_name,
+                        base_model,
+                        self.max_lora_rank,
+                        0,
+                    ),
+                    dtype=self.dtype,
+                    device=device,
+                )
+
+        if self.lora_added_tokens_size > 0:
+            self.new_embeddings_buffer["input_embeddings"] = torch.empty(
+                (
+                    self.max_loras_per_batch,
+                    self.lora_added_tokens_size,
+                    self.embedding_dim,
+                ),
+                dtype=self.dtype,
+                device=device,
+            )
+
+        if "embed_tokens" in self.target_modules:
+            init_embedding_buffer(
+                self.embedding_A_buffer,
+                self.target_modules,
+                self.get_embedding_lora_A_shape,
+            )
+
+            init_embedding_buffer(
+                self.embedding_B_buffer,
+                self.target_modules,
+                self.get_embedding_lora_B_shape,
+            )
+
+        if "lm_head" in self.target_modules:
+            init_embedding_buffer(
+                self.lm_head_A_buffer,
+                self.target_modules,
+                self.get_embedding_lora_A_shape,
+            )
+
+            init_embedding_buffer(
+                self.lm_head_B_buffer,
+                self.target_modules,
+                self.get_embedding_lora_B_shape,
+            )
+
         init_buffer(
             self.A_buffer,
             self.target_modules,
@@ -192,6 +296,8 @@ class LoRAMemoryPool:
         lora_adapters: Dict[str, LoRAAdapter],
         lora_modules: List[Dict[str, BaseLayerWithLoRA]],
         lora_refs: Dict[str, LoRARef],
+        lora_embed_tokens_module: Dict[str, BaseLayerWithLoRA],
+        lora_lm_head_module: Dict[str, BaseLayerWithLoRA],
     ):
         def get_available_buffer_slot():
             # 1. Prioritize empty slots
@@ -244,7 +350,12 @@ class LoRAMemoryPool:
                 buffer_id = get_available_buffer_slot()
                 lora_adapter = lora_adapters.get(uid, None)
                 self.load_lora_weight_to_buffer(
-                    uid, buffer_id, lora_adapter, lora_modules
+                    uid,
+                    buffer_id,
+                    lora_adapter,
+                    lora_modules,
+                    lora_embed_tokens_module,
+                    lora_lm_head_module,
                 )
                 self.uid_to_buffer_id[uid] = buffer_id
                 self.buffer_id_to_uid[buffer_id] = uid
@@ -255,6 +366,8 @@ class LoRAMemoryPool:
         buffer_id: int,
         lora_adapter: LoRAAdapter,
         lora_modules: List[Dict[str, BaseLayerWithLoRA]],
+        lora_embed_tokens_module: Dict[str, BaseLayerWithLoRA],
+        lora_lm_head_module: Dict[str, BaseLayerWithLoRA],
     ):
         def load_lora_weight_tensor(
             buffer_view: torch.Tensor, weight: Optional[torch.Tensor]
@@ -273,6 +386,12 @@ class LoRAMemoryPool:
             for i in range(self.num_layer):
                 for k in self.A_buffer.keys():
                     self.A_buffer[k][i][buffer_id] = 0
+
+            for k in self.embedding_A_buffer.keys():
+                self.embedding_A_buffer[k][buffer_id] = 0
+
+            for k in self.lm_head_A_buffer.keys():
+                self.lm_head_A_buffer[k][buffer_id] = 0
             return
 
         assert lora_adapter is not None
@@ -321,9 +440,126 @@ class LoRAMemoryPool:
                 buffer_view = target_buffer[buffer_id, :, :lora_rank]
                 load_lora_weight_tensor(buffer_view, weights)
 
+        if lora_adapter.embedding_layers:
+
+            org_vocab_size = self.base_hf_config.vocab_size
+            lora_added_tokens_size = lora_adapter.config.lora_added_tokens_size
+            # Only when LoRA is applied to the embedding layer will it have the extra-token issue that needs to be resolved.
+            # Load embeddings weights for extra tokens to buffer
+            if lora_adapter.added_tokens_embeddings:
+                for name, weights in lora_adapter.added_tokens_embeddings.items():
+                    if "input_embeddings" in name:
+                        buffer_view = self.new_embeddings_buffer["input_embeddings"][
+                            buffer_id, :lora_added_tokens_size
+                        ]
+                        load_lora_weight_tensor(buffer_view, weights)
+
+            # load vocab_emb and lm_head
+            for name, weights in lora_adapter.embedding_layers.items():
+                target_module = get_target_module_name(name, self.target_modules)
+                if (
+                    target_module == "embed_tokens"
+                    and "embed_tokens" in name
+                    and ("lora_embedding_A" in name or "lora_A" in name)
+                ):
+                    buffer_view = self.embedding_A_buffer[target_module][
+                        buffer_id,
+                        :lora_rank,
+                        : (org_vocab_size + lora_added_tokens_size),
+                    ]
+                    load_lora_weight_tensor(buffer_view, weights)
+                elif (
+                    target_module == "embed_tokens"
+                    and "embed_tokens" in name
+                    and ("lora_embedding_B" in name or "lora_B" in name)
+                ):
+                    lora_b_weights = weights
+                    # [to-do] support TP
+                    # if self.tp_size > 1:
+                    #     cur_module = lora_embeddings_modules[target_module]
+                    #     for module_name, module in cur_module:
+                    #         lora_b_weights = module.slice_lora_b_weights(
+                    #             lora_b_weights, self.tp_rank
+                    #         )
+
+                    buffer_view = self.embedding_B_buffer[target_module][
+                        buffer_id, :, :lora_rank
+                    ]
+                    load_lora_weight_tensor(buffer_view, lora_b_weights)
+
+                elif (
+                    target_module == "lm_head"
+                    and "lm_head" in name
+                    and ("lora_embedding_A" in name or "lora_A" in name)
+                ):
+                    buffer_view = self.lm_head_A_buffer[target_module][
+                        # buffer_id, :, :lora_rank
+                        buffer_id,
+                        :lora_rank,
+                        :,
+                    ]
+                    load_lora_weight_tensor(buffer_view, weights)
+                elif (
+                    target_module == "lm_head"
+                    and "lm_head" in name
+                    and ("lora_embedding_B" in name or "lora_B" in name)
+                ):
+                    lora_b_weights = weights
+                    # [to-do] support TP
+                    # if self.tp_size > 1:
+                    #     cur_module = lora_embeddings_modules[target_module]
+                    #     for module_name, module in cur_module:
+                    #         lora_b_weights = module.slice_lora_b_weights(
+                    #             lora_b_weights, self.tp_rank
+                    #         )
+
+                    buffer_view = self.lm_head_B_buffer[target_module][
+                        # buffer_id, :lora_rank, : org_vocab_size + extra_vocab_size
+                        buffer_id,
+                        : (org_vocab_size + self.lora_added_tokens_size),
+                        :lora_rank,
+                    ]
+                    load_lora_weight_tensor(buffer_view, lora_b_weights)
+
+    def get_embedding_tensor(
+        self, target_module: str, lora_type: LoRAType
+    ) -> Optional[torch.Tensor]:
+        """
+        Get LoRA tensor for non-layer modules (embed_tokens, lm_head).
+
+        Args:
+            target_module: Module name, either "embed_tokens" or "lm_head"
+            lora_type: Either LoRAType.LORA_A or LoRAType.LORA_B
+
+        Returns:
+            The corresponding buffer tensor, or None if not available
+        """
+
+        if target_module == "added_tokens":
+            if (
+                self.lora_added_tokens_size is not None
+                and self.lora_added_tokens_size > 0
+            ):
+                return self.new_embeddings_buffer["input_embeddings"]
+            return None
+        elif target_module == "embed_tokens":
+            if lora_type == LoRAType.LORA_A:
+                return self.embedding_A_buffer[target_module]
+            return self.embedding_B_buffer[target_module]
+        elif target_module == "lm_head":
+            if lora_type == LoRAType.LORA_A:
+                return self.lm_head_A_buffer[target_module]
+            return self.lm_head_B_buffer[target_module]
+
+        raise ValueError(
+            f"Invalid target_module '{target_module}'. "
+            f"Expected 'embed_tokens' or 'lm_head'."
+        )
+
     def get_tensor(
         self, target_module: str, layer_id: int, lora_type: LoRAType
     ) -> torch.Tensor:
+
         if lora_type == LoRAType.LORA_A:
             return self.A_buffer[target_module][layer_id]
 

python/sglang/srt/lora/triton_ops/__init__.py

@@ -1,5 +1,6 @@
 from .chunked_sgmv_expand import chunked_sgmv_lora_expand_forward
 from .chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward
+from .embedding_lora_a import embedding_lora_a_fwd
 from .gate_up_lora_b import gate_up_lora_b_fwd
 from .qkv_lora_b import qkv_lora_b_fwd
 from .sgemm_lora_a import sgemm_lora_a_fwd
@@ -12,4 +13,5 @@ __all__ = [
     "sgemm_lora_b_fwd",
     "chunked_sgmv_lora_shrink_forward",
     "chunked_sgmv_lora_expand_forward",
+    "embedding_lora_a_fwd",
 ]

python/sglang/srt/lora/triton_ops/embedding_lora_a.py

@@ -0,0 +1,186 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _embedding_lora_a_kernel(
+    # Pointers to tensors
+    input_ids,
+    weights,
+    output,
+    extra_embeddings,
+    # Dimensions
+    vocab_size,
+    rank,
+    num_loras,
+    # Strides
+    w_stride_0,  # stride for lora index
+    w_stride_1,  # stride for rank
+    w_stride_2,  # stride for vocab
+    output_stride_0,
+    output_stride_1,
+    extra_emb_stride_0,  # stride for lora index
+    extra_emb_stride_1,  # stride for token
+    extra_emb_stride_2,  # stride for hidden dim (= rank for extra embeddings)
+    # Batch info
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    # Meta-parameters
+    BLOCK_RANK: tl.constexpr,
+    HAS_EXTRA_EMBEDDINGS: tl.constexpr,
+):
+    """
+    Embedding lookup for LoRA A weights with support for extra tokens.
+
+    Each program handles one token across a block of rank dimensions.
+    Grid: (cdiv(max_len, 1), bs) - one program per token in each batch
+    """
+    batch_id = tl.program_id(axis=1)
+    token_idx = tl.program_id(axis=0)
+
+    w_index = tl.load(weight_indices + batch_id)
+    rank_val = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, skip
+    if rank_val == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_len = tl.load(seg_lens + batch_id)
+
+    # Check if this token is within the segment
+    if token_idx >= seg_len:
+        return
+
+    # Load the token ID
+    token_id = tl.load(input_ids + seg_start + token_idx)
+
+    # Process in chunks of BLOCK_RANK dimensions
+    num_blocks = tl.cdiv(rank_val, BLOCK_RANK)
+
+    for block_id in range(num_blocks):
+        rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK
+        rank_mask = rank_offset < rank_val
+
+        # Check if this is an extra token
+        is_extra_token = token_id >= vocab_size
+
+        if HAS_EXTRA_EMBEDDINGS and is_extra_token:
+            # Use extra embeddings
+            extra_token_id = token_id - vocab_size
+            extra_emb_ptr = (
+                extra_embeddings
+                + w_index * extra_emb_stride_0
+                + extra_token_id * extra_emb_stride_1
+                + rank_offset * extra_emb_stride_2
+            )
+            emb_values = tl.load(extra_emb_ptr, mask=rank_mask, other=0.0)
+        else:
+            # Use regular LoRA A weights
+            # weights shape: (num_loras, rank, vocab_size)
+            # We need to load weights[w_index, rank_offset, token_id]
+            token_id_clamped = tl.minimum(token_id, vocab_size - 1)
+            weight_ptr = (
+                weights
+                + w_index * w_stride_0
+                + rank_offset * w_stride_1
+                + token_id_clamped * w_stride_2
+            )
+            emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0)
+
+        # Write to output
+        output_ptr = (
+            output
+            + (seg_start + token_idx) * output_stride_0
+            + rank_offset * output_stride_1
+        )
+        tl.store(output_ptr, emb_values, mask=rank_mask)
+
+
+def embedding_lora_a_fwd(
+    input_ids: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    vocab_size: int,
+    extra_embeddings: torch.Tensor = None,
+) -> torch.Tensor:
+    """
+    Forward pass for LoRA A embedding lookup.
+
+    Args:
+        input_ids: (s,) token IDs
+        weights: (num_loras, rank, vocab_size) LoRA A embedding weights
+        batch_info: LoRABatchInfo containing batch information
+        vocab_size: base vocabulary size
+        extra_embeddings: (num_loras, num_extra_tokens, rank) extra token embeddings
+
+    Returns:
+        output: (s, rank) embedded features
+    """
+    assert input_ids.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(input_ids.shape) == 1
+    assert len(weights.shape) == 3
+
+    S = input_ids.shape[0]
+    num_loras = weights.shape[0]
+    rank = weights.shape[1]
+    vocab_size_weights = weights.shape[2]
+
+    # Block size for rank dimension
+    BLOCK_RANK = 128
+
+    has_extra_embeddings = extra_embeddings is not None
+
+    if has_extra_embeddings:
+        assert extra_embeddings.is_contiguous()
+        extra_emb_stride = (
+            extra_embeddings.stride(0),
+            extra_embeddings.stride(1),
+            extra_embeddings.stride(2),
+        )
+    else:
+        # Create dummy tensor to satisfy Triton
+        extra_embeddings = torch.empty(
+            (1, 1, 1), device=input_ids.device, dtype=weights.dtype
+        )
+        extra_emb_stride = (1, 1, 1)
+
+    # Grid: one program per token in each batch segment
+    grid = (
+        batch_info.max_len,
+        batch_info.bs,
+    )
+
+    output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
+
+    _embedding_lora_a_kernel[grid](
+        input_ids,
+        weights,
+        output,
+        extra_embeddings,
+        vocab_size,
+        rank,
+        num_loras,
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        extra_emb_stride[0],
+        extra_emb_stride[1],
+        extra_emb_stride[2],
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        BLOCK_RANK,
+        has_extra_embeddings,
+    )
+
+    return output

python/sglang/srt/lora/utils.py

@@ -46,7 +46,11 @@ class LoRAType(Enum):
 
 
 def get_hidden_dim(
-    module_name: str, config: AutoConfig, base_model: torch.nn.Module, layer_idx: int
+    module_name: str,
+    config: AutoConfig,
+    base_model: torch.nn.Module,
+    layer_idx: int,
+    lora_added_vocab_size: int = 0,
 ) -> Tuple[int]:
     """
     Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
@@ -78,6 +82,14 @@ def get_hidden_dim(
             return config.hidden_size, config.intermediate_size * 2
         elif module_name == "down_proj":
             return config.intermediate_size, config.hidden_size
+        elif module_name == "embed_tokens":
+            # For embedding: input is vocab_size (as embedding lookup), output is hidden_size
+            # if contain extra tokens will be added; otherwise is 0.
+            return config.vocab_size + lora_added_vocab_size, config.hidden_size
+        elif module_name == "lm_head":
+            # For lm_head: input is hidden_size, output is vocab_size
+            # if contain extra tokens will be added; otherwise is 0.
+            return config.hidden_size, config.vocab_size + lora_added_vocab_size
         else:
             raise NotImplementedError()
 
@@ -95,6 +107,12 @@ def get_normalized_target_modules(
         "v_proj": "qkv_proj",
         "gate_proj": "gate_up_proj",
         "up_proj": "gate_up_proj",
+        "embed_tokens": "embed_tokens",
+        "vocab_emb": "embed_tokens",
+        "embeddings": "embed_tokens",
+        "word_embeddings": "embed_tokens",
+        "lm_head": "lm_head",
+        "output": "lm_head",
     }
 
     result = set()
@@ -131,4 +149,5 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
     )
 
 
+EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
 ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj"]

python/sglang/srt/utils/common.py

@@ -3322,6 +3322,8 @@ SUPPORTED_LORA_TARGET_MODULES = [
     "down_proj",
     "qkv_proj",
     "gate_up_proj",
+    "embed_tokens",
+    "lm_head",
 ]
 
 LORA_TARGET_ALL_MODULES = "all"

python/sglang/test/runners.py

@@ -436,6 +436,7 @@ class HFRunner:
                 )
             else:
                 model = base_model
+
             if patch_model_do_sample_false:
                 model.generation_config.do_sample = False
             outputs = model.generate(
@@ -455,6 +456,7 @@ class HFRunner:
             text = tokenizer.decode(
                 outputs[0][0][len(input_ids[0]) :], skip_special_tokens=True
             )
+
             # Check if the text is empty or only whitespace.
             if not text.strip():
                 raise ValueError(