Class LoraManager

tensorrt_llm/lora_manager.py:640–1269

class LoraManager(object):
    LORA_MODULE_IDS = {
        "attn_qkv": 0,
        "attn_q": 1,
        "attn_k": 2,
        "attn_v": 3,
        "attn_dense": 4,
        "mlp_h_to_4h": 5,
        "mlp_4h_to_h": 6,
        "mlp_gate": 7,
        "cross_attn_qkv": 8,
        "cross_attn_q": 9,
        "cross_attn_k": 10,
        "cross_attn_v": 11,
        "cross_attn_dense": 12,
        "moe_h_to_4h": 13,
        "moe_4h_to_h": 14,
        "moe_gate": 15,
        "moe_router": 16,
        "mlp_router": 17,
        "mlp_gate_up": 18,
    }

    def __init__(
        self,
        *,
        mapping: Mapping,
        model_config: "ModelConfig",
        cpp_peft_cache_manager: tb_internal.batch_manager.PeftCacheManager | None = None,
    ):
        """Constructor.

        Args:
            mapping (Mapping): Parallelism related information.
            model_config (ModelConfig): model configuration python class instance.
            cpp_peft_cache_manager (PeftCacheManager, optional): used by is_adapter_in_cpu_cache method, that's used for
                a performance optimization with LoRA of not sending the LoRA adapter weights with every LLM request when
                the adapter is already loaded in the LoRA CPU cache.
        """
        # _lora_uid_to_low_ranks: dict[str -> dict[int -> dict[str -> int]]]
        # {
        #     uid: {
        #         0: {
        #             lora_module: int
        #         }, # layer_0_rank,
        #         1: {
        #             lora_module: int
        #         }, # layer_1_rank,
        #         ...
        #     }
        # }

        # _lora_weights_pointers_list: dict[str -> dict[int -> dict[str -> [Tensor, Tensor]]]]
        # {
        #     uid: {
        #         0: {
        #             lora_module: [t_in, t_out]
        #         }, # layer_0,