sgl-project--sglang/python/sglang/srt/lora/mem_pool.py

import logging
import re
from typing import (
    Callable,
    Dict,
    Iterable,
    Iterator,
    List,
    Optional,
    Set,
    Tuple,
    Union,
    overload,
)

import torch

from sglang.srt.distributed import (
    divide,
    get_pp_group,
)
from sglang.srt.environ import envs
from sglang.srt.lora.eviction_policy import get_eviction_policy
from sglang.srt.lora.layers import BaseLayerWithLoRA
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,
    REPLICATED_LINEAR_LORA_NAMES,
    ROW_PARALLELISM_LINEAR_LORA_NAMES,
    LoRAType,
    copy_weight_into_buffer,
    get_hidden_dim,
    get_lm_head_lora_b_shard_size,
    get_normalized_target_modules,
    get_stacked_multiply,
    get_target_module_name,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import is_pin_memory_available
from sglang.srt.utils.hf_transformers_utils import AutoConfig

_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()

logger = logging.getLogger(__name__)


class EmptySlot:
    """
    Singleton class to represent an empty slot in the memory pool.
    This is used to improve readability by not using special str as a placeholder.
    """

    __slots__ = ()

    def __repr__(self):
        return "|EMPTY|"

    def __new__(cls):
        if not hasattr(cls, "_instance"):
            cls._instance = super().__new__(cls)
        return cls._instance


EMPTY_SLOT = EmptySlot()


@overload
def append_cache_key_suffix(cache_keys: str, suffix: str) -> str: ...


@overload
def append_cache_key_suffix(
    cache_keys: Dict[int, str], suffix: str
) -> Dict[int, str]: ...


def append_cache_key_suffix(
    cache_keys: Union[str, Dict[int, str]],
    suffix: str,
) -> Union[str, Dict[int, str]]:
    if isinstance(cache_keys, dict):
        return {
            expert_id: f"{cache_key}#{suffix}"
            for expert_id, cache_key in cache_keys.items()
        }
    return f"{cache_keys}#{suffix}"


def _get_moe_ep_context() -> Tuple[int, int]:
    """Return `(moe_ep_size, moe_ep_rank)`, or `(1, 0)` if the MoE EP group
    is not initialized (hermetic tests or pure-TP launches)."""
    try:
        return get_parallel().moe_ep_size, get_parallel().moe_ep_rank
    except Exception:  # pragma: no cover - MoE EP group not initialized
        return 1, 0


def _get_moe_tp_context() -> Tuple[int, int]:
    """Return `(moe_tp_size, moe_tp_rank)`, or `(1, 0)` if the MoE TP group
    is not initialized. Under `--tp N --ep N` the outer attention TP group
    is consumed entirely by EP, leaving `moe_tp_size == 1`, so per-expert
    MoE weights are NOT sharded along their inner dim even though attention
    weights are."""
    try:
        return get_parallel().moe_tp_size, get_parallel().moe_tp_rank
    except Exception:  # pragma: no cover - MoE TP group not initialized
        return 1, 0


def _moe_runner_keeps_global_expert_ids() -> bool:
    """True if the active MoE runner keeps global `topk_ids` instead of
    remapping to local IDs. Mirrors the predicate in `StandardDispatcher`."""
    try:
        from sglang.srt.layers.moe.utils import get_moe_runner_backend

        b = get_moe_runner_backend()
        return (
            b.is_flashinfer_cutlass()
            or b.is_flashinfer_cutedsl()
            or b.is_experimental_sgl_trtllm()
            or b.is_flashinfer_trtllm_routed()
        )
    except Exception:  # pragma: no cover - backend not initialized
        return False


class LoRAMemoryPool:
    """Class for memory pool management of lora modules"""

    def __init__(
        self,
        base_hf_config: AutoConfig,
        max_loras_per_batch: int,
        dtype: torch.dtype,
        tp_size: int,
        tp_rank: int,
        max_lora_rank: int,
        target_modules: Set[str],
        base_model: torch.nn.Module,
        eviction_policy: str,
        lora_added_tokens_size: int,
        experts_shared_outer_loras: bool = False,
        strict_loading: bool = False,
        enable_lora_overlap_loading: bool = False,
    ):
        self.base_hf_config: AutoConfig = base_hf_config
        self.num_layer: int = base_hf_config.num_hidden_layers
        self.max_loras_per_batch: int = max_loras_per_batch
        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
        self.experts_shared_outer_loras: bool = experts_shared_outer_loras
        self.strict_loading: bool = strict_loading
        self.enable_lora_overlap_loading: bool = enable_lora_overlap_loading
        self.pin_memory_available: bool = is_pin_memory_available()

        # Under EP with a Triton/DeepGEMM runner, `StandardDispatcher` remaps
        # global `topk_ids` -> local expert IDs before the MoE kernel, so
        # per-expert LoRA buffers must be sized and keyed by the local slice.
        # FlashInfer CUTLASS/CuteDSL/TRTLLM-routed keep global IDs, and an
        # uneven expert split (`num_experts % moe_ep_size != 0`, shouldn't
        # happen in practice) is also treated as globally-keyed so we don't
        # silently truncate experts.
        self.moe_ep_size, self.moe_ep_rank = _get_moe_ep_context()
        num_experts_global = self._get_num_experts(base_model)
        self.moe_use_local_expert_ids = (
            self.moe_ep_size > 1
            and not _moe_runner_keeps_global_expert_ids()
            and num_experts_global % self.moe_ep_size == 0
        )

        # Per-expert MoE weights are sharded by `moe_tp_size`, NOT the outer
        # `tp_size`: `moe_tp_size = tp_size // ep_size // dp_size`, so under
        # e.g. `--tp 4 --ep 4` each rank holds full-width expert weights
        # (`moe_tp_size == 1`). Sizing per-expert LoRA buffers by `tp_size`
        # here would yield a 4x-narrower inner dim than the adapter weight
        # (which `FusedMoEWithLoRA.slice_moe_lora_{a,b}_weights` correctly
        # skip-slices when `moe_tp_size <= 1`), producing a shape-mismatch
        # assert during weight load. Non-MoE modules still shard by
        # `tp_size` because attention TP is unchanged.
        self.moe_tp_size, self.moe_tp_rank = _get_moe_tp_context()

        # Initialize eviction policy
        self.eviction_policy = get_eviction_policy(eviction_policy)

        # Both A_buffer and B_buffer maps lora weight names to its buffer space.
        # Standard LoRA (3D): [num_loras, rank, hidden_dim]
        # MoE LoRA (4D): [num_loras, num_experts, rank, hidden_dim]
        # The dimensionality is determined by the module type (MoE vs standard)
        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] = {}

        # Buffer idx -> lora uid in memory pool
        # All uids are initialized as `EmptySlot` for empty buffer slots
        # Here we don't initialize to None since None is a valid uid
        self.buffer_id_to_uid: List[Union[str, None, EmptySlot]] = [
            EMPTY_SLOT
        ] * self.max_loras_per_batch

        # Cache lm_head shard_indices from the base model so that buffer
        # allocation uses the same sharding as the base ParallelLMHead layer.
        self.lm_head_shard_indices = None
        if "lm_head" in target_modules and tp_size > 1:
            from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead

            for _, module in base_model.named_modules():
                if isinstance(module, ParallelLMHead):
                    self.lm_head_shard_indices = module.shard_indices
                    break

        self.init_buffers(base_model)

    def can_support(self, config: Union[LoRAConfig, Iterable[LoRAConfig]]) -> bool:
        """
        Check if the memory pool can support the given LoRA adapters.
        """

        def _can_support(config: LoRAConfig) -> bool:
            """
            Check if the memory pool can support a single LoRA adapter.
            """
            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)
            if "all" in target_module_names:
                return True
            return target_module_names.issubset(self.target_modules)

        if isinstance(config, LoRAConfig):
            return _can_support(config)
        else:
            return all(_can_support(x) for x in config)

    def is_moe_module(self, module_name: str) -> bool:
        """Check if module is part of MoE experts."""
        return "moe" in module_name

    @staticmethod
    def _get_num_experts(base_model: torch.nn.Module) -> int:
        cfg = base_model.config
        if hasattr(cfg, "get_text_config"):
            cfg = cfg.get_text_config()
        return (
            getattr(cfg, "num_experts", None)
            or getattr(cfg, "num_local_experts", None)
            or getattr(cfg, "n_routed_experts", None)
            or 1
        )

    @staticmethod
    def _has_moe_module(base_model: torch.nn.Module) -> bool:
        # Config-only detection isn't reliable: some dense configs (e.g.
        # `Qwen3_5TextConfig`) inherit `num_experts > 1` from an MoE parent.
        # Walk the loaded model for an actual FusedMoE instance before we
        # commit to allocating 4D per-expert LoRA buffers.
        from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE

        return any(isinstance(m, FusedMoE) for m in base_model.modules())

    def _get_num_local_experts(self, base_model: torch.nn.Module) -> int:
        """Experts owned by this rank. Equals the global count when EP is
        off, the runner keeps global IDs, or the split isn't even (all
        three cases fold into `moe_use_local_expert_ids == False`)."""
        total = self._get_num_experts(base_model)
        if not self.moe_use_local_expert_ids:
            return total
        return total // self.moe_ep_size

    def _global_to_local_expert_id(self, global_eid: int) -> Optional[int]:
        """Map a global expert id to this rank's local id, or `None` if
        the expert is not owned by this rank. Pass-through when buffers
        are globally-keyed."""
        if not self.moe_use_local_expert_ids:
            return global_eid
        local = global_eid - self.moe_ep_rank * self._num_experts_local
        return local if 0 <= local < self._num_experts_local else None

    def _iter_local_expert_weights(
        self,
        weights: Union[torch.Tensor, Dict[int, torch.Tensor]],
        cache_keys: Union[str, Dict[int, str]],
    ) -> Iterator[Tuple[int, torch.Tensor, str]]:
        """Yield `(local_expert_id, weight, cache_key)` triples for per-expert
        MoE LoRA inputs, filtered/remapped to this rank's slice. Accepts either
        a `{global_eid: 2D tensor}` dict or a 3D `[num_experts, *, *]` tensor."""
        if isinstance(weights, dict):
            assert isinstance(cache_keys, dict)
            for gid, w in weights.items():
                lid = self._global_to_local_expert_id(gid)
                if lid is not None:
                    yield lid, w, cache_keys[gid]
            return

        if isinstance(weights, torch.Tensor) and weights.dim() == 3:
            assert isinstance(cache_keys, str)
            total = weights.shape[0]
            if self.moe_use_local_expert_ids:
                start = self.moe_ep_rank * self._num_experts_local
                count = max(0, min(self._num_experts_local, total - start))
            else:
                start, count = 0, total
            for i in range(count):
                yield (
                    i,
                    weights[start + i],
                    append_cache_key_suffix(cache_keys, f"expert{start + i}"),
                )
            return

        raise TypeError(
            f"Expected dict or 3D torch.Tensor, got {type(weights).__name__}."
        )

    def _get_standard_shape(
        self,
        module_name: str,
        base_model: torch.nn.Module,
        max_lora_dim: int,
        layer_idx: int,
    ) -> Tuple[int]:
        """Get 3D shape for standard (non-MoE) modules."""
        input_dim, _ = get_hidden_dim(
            module_name, self.base_hf_config, base_model, layer_idx
        )
        c = get_stacked_multiply(module_name, base_model)
        if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
            input_dim = divide(input_dim, self.tp_size)
        return (self.max_loras_per_batch, max_lora_dim * c, input_dim)

    def get_lora_A_shape(
        self,
        module_name: str,
        base_model: torch.nn.Module,
        max_lora_dim: int,
        layer_idx: int,
    ) -> Tuple[int]:
        """
        Get shape for LoRA A weights. Automatically returns 3D or 4D based on module type.

        Returns:
            - Standard: [num_loras, rank, hidden_dim]
            - MoE: [num_loras, num_experts, rank, hidden_dim]
        """
        input_dim, _ = get_hidden_dim(
            module_name, self.base_hf_config, base_model, layer_idx
        )
        c = get_stacked_multiply(module_name, base_model)
        # MoE modules shard along `moe_tp_size`, not the outer `tp_size`.
        effective_tp_size = (
            self.moe_tp_size if self.is_moe_module(module_name) else self.tp_size
        )
        if (
            effective_tp_size > 1
            and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES
            and module_name not in REPLICATED_LINEAR_LORA_NAMES
        ):
            input_dim = divide(input_dim, effective_tp_size)

        if self.is_moe_module(module_name):
            expert_dim = self._get_num_local_experts(base_model)
            if self.experts_shared_outer_loras and module_name == "gate_up_proj_moe":
                expert_dim = 1
            return (
                self.max_loras_per_batch,
                expert_dim,
                max_lora_dim * c,
                input_dim,
            )
        else:
            return (self.max_loras_per_batch, max_lora_dim * c, 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
        )
        # Embedding LoRA A is kept unsharded (full vocab) across TP ranks.
        # Each rank does a full lookup; no vocab-dimension splitting needed.
        return (
            self.max_loras_per_batch,
            max_lora_dim,
            input_dim,
        )

    def _column_parallel_lora_b_per_rank_dim(
        self,
        module_name: str,
        total_output_dim: int,
        effective_tp_size: int,
    ) -> int:
        """Per-rank LoRA B output dim for column-parallel modules.

        For most modules this is just an even split. For ``qkv_proj`` when
        ``effective_tp_size > num_key_value_heads``, the underlying
        :class:`QKVParallelLinear` *replicates* each KV head across
        ``tp_size // num_kv_heads`` ranks instead of dividing further, so
        each rank owns ``head_dim`` of K/V (not ``head_dim * num_kv_heads
        / tp_size``). A naive ``divide(total, tp_size)`` undersizes the
        buffer and produces a shape mismatch when the
        :meth:`QKVParallelLinearWithLoRA.slice_lora_b_weights` slice runs.
        """
        if module_name != "qkv_proj":
            return divide(total_output_dim, effective_tp_size)

        cfg = self.base_hf_config
        if hasattr(cfg, "get_text_config"):
            cfg = cfg.get_text_config()
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_kv_heads is None or num_kv_heads >= effective_tp_size:
            return divide(total_output_dim, effective_tp_size)

        head_dim = getattr(cfg, "head_dim", None) or (
            cfg.hidden_size // cfg.num_attention_heads
        )
        kv_dim_total = 2 * num_kv_heads * head_dim
        q_dim_total = total_output_dim - kv_dim_total
        q_per_rank = divide(q_dim_total, effective_tp_size)
        return q_per_rank + 2 * head_dim

    def get_lora_B_shape(
        self,
        module_name: str,
        base_model: torch.nn.Module,
        max_lora_dim: int,
        layer_idx: int,
    ) -> Tuple[int]:
        """
        Get shape for LoRA B weights. Automatically returns 3D or 4D based on module type.

        Returns:
            - Standard: [num_loras, output_dim, rank]
            - MoE: [num_loras, num_experts, output_dim, rank]
        """
        _, output_dim = get_hidden_dim(
            module_name, self.base_hf_config, base_model, layer_idx
        )
        # MoE modules shard along `moe_tp_size`, not the outer `tp_size`.
        effective_tp_size = (
            self.moe_tp_size if self.is_moe_module(module_name) else self.tp_size
        )
        if (
            effective_tp_size > 1
            and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES
            and module_name not in REPLICATED_LINEAR_LORA_NAMES
        ):
            output_dim = self._column_parallel_lora_b_per_rank_dim(
                module_name, output_dim, effective_tp_size
            )

        # Check if MoE module and return appropriate shape
        if self.is_moe_module(module_name):
            expert_dim = self._get_num_local_experts(base_model)
            if self.experts_shared_outer_loras and module_name == "down_proj_moe":
                expert_dim = 1
            return (self.max_loras_per_batch, expert_dim, output_dim, max_lora_dim)
        else:
            return (self.max_loras_per_batch, output_dim, 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
        )
        # lm_head is column-parallel so B is sharded; embed_tokens B stays
        # unsharded (base output is all-reduced to full embed_dim).
        if module_name == "lm_head":
            output_dim = get_lm_head_lora_b_shard_size(
                output_dim,
                shard_indices=self.lm_head_shard_indices,
            )
        return (
            self.max_loras_per_batch,
            output_dim,
            max_lora_dim,
        )

    def init_buffers(self, base_model: torch.nn.Module):
        self.base_model = base_model
        device = next(base_model.parameters()).device

        # Cached once so the per-expert load path doesn't re-walk the HF
        # config for every adapter.
        self._num_experts_local: int = self._get_num_local_experts(base_model)

        def init_buffer(
            buffer: Dict[str, List[torch.Tensor]],
            target_modules: Set[str],
            get_lora_shape_fn: Callable[[str, torch.nn.Module, int, int], Tuple[int]],
        ):
            cfg = base_model.config
            if hasattr(cfg, "get_text_config"):
                cfg = cfg.get_text_config()
            has_shared_experts = (
                hasattr(cfg, "shared_expert_intermediate_size")
                and cfg.shared_expert_intermediate_size > 0
            ) or (getattr(cfg, "n_shared_experts", 0) or 0) > 0
            has_moe = self._has_moe_module(base_model)

            # Shape functions automatically handle both 3D (standard) and 4D (MoE)
            target_modules = target_modules - set(EMBEDDING_NAMES)
            for module_name in target_modules:
                # Special handling for ambiguous target modules that can be in different contexts
                ambiguous_modules = {"gate_up_proj", "down_proj"}
                if module_name in ambiguous_modules and has_moe:
                    # Allocate shared expert version (3D) only when model has shared experts
                    if has_shared_experts:
                        buffer[module_name] = [
                            torch.zeros(
                                get_lora_shape_fn(
                                    module_name, base_model, self.max_lora_rank, idx
                                ),
                                dtype=self.dtype,
                                device=device,
                            )
                            for idx in range(self.num_layer)
                        ]

                    # MoE expert version (4D)
                    moe_key = f"{module_name}_moe"
                    buffer[moe_key] = [
                        torch.zeros(
                            get_lora_shape_fn(
                                moe_key, base_model, self.max_lora_rank, idx
                            ),
                            dtype=self.dtype,
                            device=device,
                        )
                        for idx in range(self.num_layer)
                    ]
                else:
                    # Standard allocation for unambiguous modules
                    buffer[module_name] = [
                        torch.zeros(
                            get_lora_shape_fn(
                                module_name,
                                base_model,
                                self.max_lora_rank,
                                idx,
                            ),
                            dtype=self.dtype,
                            device=device,
                        )
                        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.zeros(
                    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.zeros(
                (
                    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,
            self.get_lora_A_shape,
        )

        init_buffer(
            self.B_buffer,
            self.target_modules,
            self.get_lora_B_shape,
        )

    def _get_maybe_cached_weight_for_transfer(
        self,
        pinned_weight_store: Dict[str, torch.Tensor],
        cache_key: str,
        weight: torch.Tensor,
    ) -> torch.Tensor:
        if (
            not self.pin_memory_available
            or weight.device.type != "cpu"
            or weight.is_pinned()
        ):
            return weight

        if not self.enable_lora_overlap_loading:
            return weight.pin_memory()

        cached_weight = pinned_weight_store.get(cache_key)
        if cached_weight is None:
            cached_weight = weight.pin_memory()
            pinned_weight_store[cache_key] = cached_weight
        elif cached_weight.shape != weight.shape or cached_weight.dtype != weight.dtype:
            raise ValueError(
                f"LoRA pinned weight cache key collision for {cache_key!r}: "
                f"cached shape={cached_weight.shape}, dtype={cached_weight.dtype}; "
                f"new shape={weight.shape}, dtype={weight.dtype}."
            )

        return cached_weight

    def prepare_lora_batch(
        self,
        cur_uids: Set[Optional[str]],
        lora_adapters: Dict[str, LoRAAdapter],
        lora_modules: List[Dict[str, BaseLayerWithLoRA]],
        lora_refs: Dict[str, LoRARef],
        lora_embed_tokens_module: Optional[BaseLayerWithLoRA],
        lora_lm_head_module: Optional[BaseLayerWithLoRA],
    ):
        def get_available_buffer_slot():
            # 1. Prioritize empty slots
            for buffer_id in range(self.max_loras_per_batch):
                if self.buffer_id_to_uid[buffer_id] == EMPTY_SLOT:
                    return buffer_id

            # 2. Memory pool is full, need to evict using policy
            candidates = set()

            for buffer_id in range(self.max_loras_per_batch):
                uid = self.buffer_id_to_uid[buffer_id]

                # Skip if this adapter is needed by current batch
                if uid in cur_uids:
                    continue

                # Skip if this adapter is pinned
                if uid is not None:
                    lora_ref = lora_refs.get(uid)
                    if lora_ref and lora_ref.pinned:
                        continue

                candidates.add(uid)

            if not candidates:
                raise ValueError(
                    "No available buffer slots found. Please ensure the number of active (pinned) loras is less than max_loras_per_batch."
                )

            # Prefer evicting LoRA adapters over the base model (None).
            # Only evict None when the batch consists entirely of LoRA requests
            # and no other adapters can be evicted.
            non_none_candidates = candidates - {None}
            if non_none_candidates:
                # Prioritize evicting actual LoRA adapters
                candidates_to_use = non_none_candidates
            else:
                # Only None is available for eviction (batch is all LoRA requests)
                candidates_to_use = candidates

            # Select victim using eviction policy
            victim_uid = self.eviction_policy.select_victim(candidates_to_use)

            # Evict the selected victim
            victim_buffer_id = self.uid_to_buffer_id[victim_uid]
            self.uid_to_buffer_id.pop(victim_uid)
            self.eviction_policy.remove(victim_uid)
            self.buffer_id_to_uid[victim_buffer_id] = EMPTY_SLOT
            logger.debug(
                f"Evicting LoRA {victim_uid} from buffer slot {victim_buffer_id}."
            )
            return victim_buffer_id

        # Mark all adapters in current batch as used (for LRU tracking)
        for uid in cur_uids:
            self.eviction_policy.mark_used(uid)

        for uid in cur_uids:
            if uid not in self.uid_to_buffer_id:
                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,
                    lora_embed_tokens_module,
                    lora_lm_head_module,
                )
                self.uid_to_buffer_id[uid] = buffer_id
                self.buffer_id_to_uid[buffer_id] = uid

    def load_lora_weight_to_buffer(
        self,
        uid: str,
        buffer_id: int,
        lora_adapter: LoRAAdapter,
        lora_modules: List[Dict[str, BaseLayerWithLoRA]],
        lora_embed_tokens_module: Optional[BaseLayerWithLoRA],
        lora_lm_head_module: Optional[BaseLayerWithLoRA],
    ):
        def load_lora_weight_tensor(
            buffer_view: torch.Tensor, weight: Optional[torch.Tensor]
        ):
            if weight is None:
                # If the particular weight is not present in the adapter, we initialize the buffer to zero
                # to avoid contamination from the residual weight of the evicted adapters.
                buffer_view.zero_()
            else:
                assert (
                    buffer_view.shape == weight.shape
                ), f"LoRA buffer shape {buffer_view.shape} does not match weight shape {weight.shape}."
                copy_weight_into_buffer(buffer_view, weight)

        if uid is None:
            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
        lora_rank = lora_adapter.config.r

        # Pre-validate weight names against target modules across all layers
        # and embedding weights.  This catches mismatches before any GPU
        # buffers are mutated.
        skipped_weight_names: set = set()
        matched_modules: set = set()
        all_weight_names: list = []
        for layer in lora_adapter.layers:
            all_weight_names.extend(layer.weights.keys())
        if lora_adapter.embedding_layers:
            all_weight_names.extend(lora_adapter.embedding_layers.keys())
        for name in all_weight_names:
            try:
                target_module = get_target_module_name(name, self.target_modules)
                matched_modules.add(target_module)
            except ValueError:
                skipped_weight_names.add(name)
        if matched_modules:
            logger.info(
                "LoRA adapter '%s': loaded weights for target modules %s.",
                uid,
                sorted(matched_modules),
            )
        if skipped_weight_names:
            msg = (
                f"LoRA adapter '{uid}': {len(skipped_weight_names)} weight(s) "
                f"skipped because they did not match any target module in "
                f"{sorted(self.target_modules)}. Skipped weights: "
                f"{sorted(skipped_weight_names)}. This likely indicates a "
                f"mismatch between the adapter's target modules and the base "
                f"model architecture."
            )
            if self.strict_loading:
                raise ValueError(msg)
            else:
                logger.warning(msg)

        for layer_id in range(self.num_layer):
            layer = lora_adapter.layers[layer_id]
            layer_weights = layer.weights
            pinned_layer_weights = layer.pinned_weights
            # - Standard: module_name -> torch.Tensor
            # - MoE: module_name -> Dict[expert_id -> torch.Tensor]
            temp_A_buffer: Dict[str, Union[torch.Tensor, Dict[int, torch.Tensor]]] = {
                target_module: None for target_module in self.A_buffer
            }
            temp_B_buffer: Dict[str, Union[torch.Tensor, Dict[int, torch.Tensor]]] = {
                target_module: None for target_module in self.B_buffer
            }
            temp_A_cache_keys: Dict[str, Optional[Union[str, Dict[int, str]]]] = {
                target_module: None for target_module in self.A_buffer
            }
            temp_B_cache_keys: Dict[str, Optional[Union[str, Dict[int, str]]]] = {
                target_module: None for target_module in self.B_buffer
            }

            for name, weights in layer_weights.items():
                target_module = get_target_module_name(name, self.target_modules)

                # Check if this is an MoE weight (has expert index in name)
                expert_match = re.search(r"experts\.(\d+)\.", name)

                if expert_match:
                    # Per-expert MoE weight — 2D tensors, one per expert
                    target_module = target_module + "_moe"
                    if temp_A_buffer[target_module] is None:
                        temp_A_buffer[target_module] = {}
                        temp_B_buffer[target_module] = {}
                        temp_A_cache_keys[target_module] = {}
                        temp_B_cache_keys[target_module] = {}

                    expert_id = int(expert_match.group(1))
                    if "lora_A" in name:
                        temp_A_buffer[target_module][expert_id] = weights
                        temp_A_cache_keys[target_module][expert_id] = name
                    else:
                        temp_B_buffer[target_module][expert_id] = weights
                        temp_B_cache_keys[target_module][expert_id] = name
                elif "experts" in name and weights.dim() == 3:
                    # Shared outer MoE weight — 3D tensor [expert_dim, rank, hidden]
                    target_module = target_module + "_moe"
                    if "lora_A" in name:
                        temp_A_buffer[target_module] = weights
                        temp_A_cache_keys[target_module] = name
                    else:
                        temp_B_buffer[target_module] = weights
                        temp_B_cache_keys[target_module] = name
                else:
                    # Standard weight — single tensor per module
                    if "lora_A" in name:
                        temp_A_buffer[target_module] = weights
                        temp_A_cache_keys[target_module] = name
                    else:
                        temp_B_buffer[target_module] = weights
                        temp_B_cache_keys[target_module] = name

            # Track which buffer keys correspond to a real wrapped module on
            # this layer. `temp_A/B_buffer` is seeded with every key in the
            # global `A/B_buffer` (union across all layer types), but a
            # hybrid-architecture layer (e.g. Qwen3.5 linear-attn vs full-attn,
            # or first-k-dense MoE) only owns a subset of those modules. The
            # buffer-copy loops below skip non-owned keys to avoid the
            # redundant zero-fills on slots no `update_lora_info` ever points
            # a forward-time module at.
            active_target_modules: Set[str] = set()
            cur_layer_modules = lora_modules[layer_id]
            for module_name, module in cur_layer_modules.items():
                # TODO (Jonahcb): check if the code can be refactored to avoid the special handling for FusedMoEWithLoRA
                # Handle FusedMoEWithLoRA specially - it contains multiple target modules
                from sglang.srt.lora.layers import FusedMoEWithLoRA

                if isinstance(module, FusedMoEWithLoRA):
                    # Per-expert MoE weights are sharded along `moe_tp_size`
                    # (= tp_size // ep_size // dp_size), so the slice index
                    # must be `moe_tp_rank`. Passing the outer `tp_rank` here
                    # produces an off-the-end slice when ep_size < tp_size
                    # (e.g. tp=4 ep=2 → ranks 2,3 slice past intermediate_size).
                    moe_target_modules = ["gate_up_proj_moe", "down_proj_moe"]
                    for target_module in moe_target_modules:
                        active_target_modules.add(target_module)
                        if temp_A_buffer.get(target_module) is not None:
                            temp_A_buffer[target_module] = (
                                module.slice_moe_lora_a_weights(
                                    temp_A_buffer[target_module],
                                    self.moe_tp_rank,
                                    target_module,
                                )
                            )
                            cache_keys = temp_A_cache_keys[target_module]
                            assert cache_keys is not None
                            temp_A_cache_keys[target_module] = append_cache_key_suffix(
                                cache_keys,
                                f"moe_tp{self.moe_tp_rank}",
                            )
                        if temp_B_buffer.get(target_module) is not None:
                            temp_B_buffer[target_module] = (
                                module.slice_moe_lora_b_weights(
                                    temp_B_buffer[target_module],
                                    self.moe_tp_rank,
                                    target_module,
                                )
                            )
                            cache_keys = temp_B_cache_keys[target_module]
                            assert cache_keys is not None
                            temp_B_cache_keys[target_module] = append_cache_key_suffix(
                                cache_keys,
                                f"moe_tp{self.moe_tp_rank}",
                            )

                    continue

                # Handle regular modules
                target_module = get_target_module_name(module_name, self.target_modules)
                # Mark active even if the adapter has no weights for this
                # module on this layer — the buffer still needs to be zeroed
                # (so a previously-evicted adapter's weights don't leak into
                # the new slot) and the wrapped layer module will read it.
                active_target_modules.add(target_module)

                if temp_A_buffer[target_module] is None:
                    # Skip weight slicing if the weight is not present in the adapter
                    continue

                # Handle standard modules
                temp_A_buffer[target_module] = module.slice_lora_a_weights(
                    temp_A_buffer[target_module], self.tp_rank
                )
                cache_keys = temp_A_cache_keys[target_module]
                assert cache_keys is not None
                temp_A_cache_keys[target_module] = append_cache_key_suffix(
                    cache_keys,
                    f"tp{self.tp_rank}",
                )

                temp_B_buffer[target_module] = module.slice_lora_b_weights(
                    temp_B_buffer[target_module], self.tp_rank
                )
                cache_keys = temp_B_cache_keys[target_module]
                assert cache_keys is not None
                temp_B_cache_keys[target_module] = append_cache_key_suffix(
                    cache_keys,
                    f"tp{self.tp_rank}",
                )

            for name, weights in temp_A_buffer.items():
                if name not in active_target_modules:
                    continue
                c = get_stacked_multiply(name, self.base_model)
                max_r = self.max_lora_rank
                target_buffer = self.A_buffer[name][layer_id]
                weights_cache_key = temp_A_cache_keys[name]

                if name in ["gate_up_proj_moe", "down_proj_moe"]:
                    if self.experts_shared_outer_loras and name == "gate_up_proj_moe":
                        if weights is None:
                            representative_weight = None
                            buffer_view = target_buffer[
                                buffer_id, 0, : lora_rank * c, :
                            ]
                            load_lora_weight_tensor(buffer_view, None)
                        elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
                            if weights.shape[0] != 1:
                                raise ValueError(
                                    f"experts_shared_outer_loras is enabled but "
                                    f"gate_up_proj_moe lora_A has expert_dim="
                                    f"{weights.shape[0]} (expected 1)."
                                )
                            assert isinstance(weights_cache_key, str)
                            weights = self._get_maybe_cached_weight_for_transfer(
                                pinned_layer_weights,
                                weights_cache_key,
                                weights,
                            )
                            representative_weight = weights[0]
                            buffer_view = target_buffer[
                                buffer_id, 0, : lora_rank * c, :
                            ]
                            load_lora_weight_tensor(buffer_view, weights[0])
                        elif isinstance(weights, dict) and len(weights) > 0:
                            if len(weights) != 1:
                                raise ValueError(
                                    f"experts_shared_outer_loras is enabled but "
                                    f"gate_up_proj_moe lora_A dict has "
                                    f"{len(weights)} entries (expected 1)."
                                )
                            rep = next(iter(weights.values()))
                            assert isinstance(weights_cache_key, dict)
                            rep_cache_key = next(iter(weights_cache_key.values()))
                            rep = self._get_maybe_cached_weight_for_transfer(
                                pinned_layer_weights, rep_cache_key, rep
                            )
                            representative_weight = rep
                            buffer_view = target_buffer[
                                buffer_id, 0, : lora_rank * c, :
                            ]
                            load_lora_weight_tensor(buffer_view, rep)
                        else:
                            raise ValueError(
                                f"Unexpected weight format for shared outer gate_up_proj_moe lora_A: "
                                f"type={type(weights)}, "
                                f"shape={weights.shape if isinstance(weights, torch.Tensor) else 'N/A'}"
                            )
                        # Place each stacked component at max_rank-spaced
                        # positions so the kernel's [:max_r] / [max_r:2*max_r]
                        # slicing is correct.
                        target_buffer[buffer_id, 0].zero_()
                        if representative_weight is not None:
                            for ci in range(c):
                                buffer_view = target_buffer[
                                    buffer_id, 0, ci * max_r : ci * max_r + lora_rank, :
                                ]
                                load_lora_weight_tensor(
                                    buffer_view,
                                    representative_weight[
                                        ci * lora_rank : (ci + 1) * lora_rank, :
                                    ],
                                )
                    elif isinstance(weights, (torch.Tensor, dict)):
                        # Zero first so any local-expert slot the adapter
                        # doesn't fill (e.g. out-of-rank under EP) is clean;
                        # then load owned slots at max_rank-spaced offsets so
                        # the MoE kernel's [:max_r] / [max_r:2*max_r] slicing
                        # is correct.
                        target_buffer[buffer_id].zero_()
                        assert isinstance(weights_cache_key, (str, dict))
                        for (
                            local_eid,
                            expert_weight,
                            expert_cache_key,
                        ) in self._iter_local_expert_weights(
                            weights, weights_cache_key
                        ):
                            if expert_weight is None:
                                continue
                            expert_weight = self._get_maybe_cached_weight_for_transfer(
                                pinned_layer_weights,
                                expert_cache_key,
                                expert_weight,
                            )
                            for ci in range(c):
                                buffer_view = target_buffer[
                                    buffer_id,
                                    local_eid,
                                    ci * max_r : ci * max_r + lora_rank,
                                    :,
                                ]
                                load_lora_weight_tensor(
                                    buffer_view,
                                    expert_weight[
                                        ci * lora_rank : (ci + 1) * lora_rank, :
                                    ],
                                )
                else:
                    buffer_view = target_buffer[buffer_id, : lora_rank * c, :]
                    if weights is not None:
                        assert isinstance(weights_cache_key, str)
                        weights = self._get_maybe_cached_weight_for_transfer(
                            pinned_layer_weights,
                            weights_cache_key,
                            weights,
                        )
                    load_lora_weight_tensor(buffer_view, weights)

            for name, weights in temp_B_buffer.items():
                if name not in active_target_modules:
                    continue
                target_buffer = self.B_buffer[name][layer_id]
                weights_cache_key = temp_B_cache_keys[name]

                if name in ["gate_up_proj_moe", "down_proj_moe"]:
                    if self.experts_shared_outer_loras and name == "down_proj_moe":
                        if weights is None:
                            buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
                            load_lora_weight_tensor(buffer_view, None)
                        elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
                            if weights.shape[0] != 1:
                                raise ValueError(
                                    f"experts_shared_outer_loras is enabled but "
                                    f"down_proj_moe lora_B has expert_dim="
                                    f"{weights.shape[0]} (expected 1)."
                                )
                            buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
                            w = weights[0]
                            assert isinstance(weights_cache_key, str)
                            if w is not None:
                                w = w * lora_adapter.scaling
                                w = self._get_maybe_cached_weight_for_transfer(
                                    pinned_layer_weights,
                                    append_cache_key_suffix(
                                        weights_cache_key, "expert0"
                                    ),
                                    w,
                                )
                            load_lora_weight_tensor(buffer_view, w)
                        elif isinstance(weights, dict) and len(weights) > 0:
                            if len(weights) != 1:
                                raise ValueError(
                                    f"experts_shared_outer_loras is enabled but "
                                    f"down_proj_moe lora_B dict has "
                                    f"{len(weights)} entries (expected 1)."
                                )
                            rep = next(iter(weights.values()))
                            assert isinstance(weights_cache_key, dict)
                            rep_cache_key = next(iter(weights_cache_key.values()))
                            buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
                            if rep is not None:
                                rep = rep * lora_adapter.scaling
                                rep = self._get_maybe_cached_weight_for_transfer(
                                    pinned_layer_weights,
                                    rep_cache_key,
                                    rep,
                                )
                            load_lora_weight_tensor(buffer_view, rep)
                        else:
                            raise ValueError(
                                f"Unexpected weight format for shared outer down_proj_moe lora_B: "
                                f"type={type(weights)}, "
                                f"shape={weights.shape if isinstance(weights, torch.Tensor) else 'N/A'}"
                            )
                        # Zero beyond loaded rank — MoE kernel reads full max_rank.
                        target_buffer[buffer_id, 0, :, lora_rank:].zero_()
                    elif isinstance(weights, (torch.Tensor, dict)):
                        # Zero out slots this rank owns but the adapter
                        # doesn't fill (padded-out / out-of-rank experts);
                        # then scale+load the ones it does.
                        target_buffer[buffer_id].zero_()
                        assert isinstance(weights_cache_key, (str, dict))
                        for (
                            local_eid,
                            w,
                            w_cache_key,
                        ) in self._iter_local_expert_weights(
                            weights, weights_cache_key
                        ):
                            if w is not None:
                                w = w * lora_adapter.scaling
                                w = self._get_maybe_cached_weight_for_transfer(
                                    pinned_layer_weights,
                                    w_cache_key,
                                    w,
                                )
                            buffer_view = target_buffer[
                                buffer_id, local_eid, :, :lora_rank
                            ]
                            load_lora_weight_tensor(buffer_view, w)
                else:
                    buffer_view = target_buffer[buffer_id, :, :lora_rank]
                    if weights is not None:
                        assert isinstance(weights_cache_key, str)
                        weights = self._get_maybe_cached_weight_for_transfer(
                            pinned_layer_weights,
                            weights_cache_key,
                            weights,
                        )
                    load_lora_weight_tensor(buffer_view, weights)
                    if _SGLANG_EXPERIMENTAL_LORA_OPTI:
                        # Zero beyond loaded rank: the experimental dense LoRA-B kernel
                        # contracts over the full padded max_rank, so the tail must be clean.
                        target_buffer[buffer_id, :, lora_rank:].zero_()

        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
            pinned_embedding_layers = lora_adapter.pinned_embedding_layers
            pinned_added_tokens_embeddings = lora_adapter.pinned_added_tokens_embeddings
            # 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
                        ]
                        weights = self._get_maybe_cached_weight_for_transfer(
                            pinned_added_tokens_embeddings,
                            name,
                            weights,
                        )
                        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),
                    ]
                    weights = self._get_maybe_cached_weight_for_transfer(
                        pinned_embedding_layers,
                        name,
                        weights,
                    )
                    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
                    # TP is supported by keeping embedding LoRA B unsharded;
                    # no slicing needed.

                    buffer_view = self.embedding_B_buffer[target_module][
                        buffer_id, :, :lora_rank
                    ]
                    lora_b_weights = self._get_maybe_cached_weight_for_transfer(
                        pinned_embedding_layers,
                        name,
                        lora_b_weights,
                    )
                    load_lora_weight_tensor(buffer_view, lora_b_weights)

                elif (
                    target_module == "lm_head"
                    and lora_lm_head_module is not None
                    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,
                        :,
                    ]
                    weights = self._get_maybe_cached_weight_for_transfer(
                        pinned_embedding_layers,
                        name,
                        weights,
                    )
                    load_lora_weight_tensor(buffer_view, weights)
                elif (
                    target_module == "lm_head"
                    and lora_lm_head_module is not None
                    and "lm_head" in name
                    and ("lora_embedding_B" in name or "lora_B" in name)
                ):
                    assert lora_lm_head_module is not None
                    lora_b_weights = weights
                    # Slice B along vocab dimension for this TP rank
                    if self.tp_size > 1:
                        lora_b_weights = lora_lm_head_module.slice_lora_b_weights(
                            lora_b_weights, self.tp_rank
                        )
                        cache_key = append_cache_key_suffix(name, f"tp{self.tp_rank}")
                    else:
                        cache_key = name

                    buffer_view = self.lm_head_B_buffer[target_module][
                        buffer_id,
                        : lora_b_weights.shape[0],
                        :lora_rank,
                    ]
                    lora_b_weights = self._get_maybe_cached_weight_for_transfer(
                        pinned_embedding_layers,
                        cache_key,
                        lora_b_weights,
                    )
                    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
                        or "lora_embedding_B" in name
                        or "lora_B" in name
                    )
                ):
                    # Only assert for genuine LoRA A/B deltas. Non-LoRA adapter
                    # entries (e.g. `base_layer.weight` emitted by PEFT for
                    # tied-embedding lm_head) fall through and are handled by
                    # the base weight loader, mirroring embed_tokens behavior.
                    # Non-last PP stages do not own lm_head, so adapters can
                    # legitimately contain lm_head LoRA weights with no local
                    # module to load them into, otherwise we should have been able to load this weight.
                    assert (
                        not get_pp_group().is_last_rank
                    ), f"Failed to load lm_head LoRA weight: {name}, this is only expected to happen on non-last PP stages."
                    continue
        else:
            # Zero out embedding/lm_head buffers for adapters without embedding LoRA
            # to avoid using garbage values from uninitialized memory
            for k in self.embedding_A_buffer.keys():
                self.embedding_A_buffer[k][buffer_id].zero_()
            for k in self.embedding_B_buffer.keys():
                self.embedding_B_buffer[k][buffer_id].zero_()
            for k in self.lm_head_A_buffer.keys():
                self.lm_head_A_buffer[k][buffer_id].zero_()
            for k in self.lm_head_B_buffer.keys():
                self.lm_head_B_buffer[k][buffer_id].zero_()
            if (
                self.lora_added_tokens_size > 0
                and "input_embeddings" in self.new_embeddings_buffer
            ):
                self.new_embeddings_buffer["input_embeddings"][buffer_id].zero_()

    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:
        """
        Get LoRA tensor buffer (automatically handles both 3D and 4D tensors).

        if lora_type == LoRAType.LORA_A:
            return self.A_buffer[target_module][layer_id]

        Args:
            target_module: Target module name (e.g., 'gate_up_proj' or 'gate_up_proj_moe' for MoE)
            layer_id: Layer index
            lora_type: LoRAType.LORA_A or LoRAType.LORA_B

        Returns:
            - 3D tensor [num_loras, rank, hidden] for standard modules
            - 4D tensor [num_loras, num_experts, rank, hidden] for MoE modules
        """
        buffer_dict = self.A_buffer if lora_type == LoRAType.LORA_A else self.B_buffer

        return buffer_dict[target_module][layer_id]

    def get_buffer_id(self, lora_uid: str):
        return self.uid_to_buffer_id[lora_uid]