Method input_buffers

tensorrt_llm/lora_manager.py:1239–1269

(self, lora_uids, mapping: Mapping, num_layers: int)

            save_val(all_configs, uid_path, "lora_config", tp_num=None, write_npy=True)

    def input_buffers(self, lora_uids, mapping: Mapping, num_layers: int):
        inputs = {}
        for layer_idx in mapping.pp_layers(num_layers):
            for lora_module in self.lora_target_modules + self.missing_qkv_modules:
                lora_ranks_ = []
                lora_ptrs_ = []
                for lora_uid in lora_uids:
                    lora_rank = 0
                    lora_ptrs = [0, 0, 0]

                    if lora_uid != "-1":
                        low_ranks = self.uid_to_low_ranks(lora_uid)

                        if (
                            layer_idx in low_ranks
                            and lora_module in low_ranks[layer_idx].keys()
                            and low_ranks[layer_idx][lora_module] != 0
                        ):
                            lora_rank = low_ranks[layer_idx][lora_module]
                            lora_ptrs = self.lora_weights_pointers_list[lora_uid][layer_idx][
                                lora_module
                            ]

                    lora_ranks_.append(lora_rank)
                    lora_ptrs_.append(lora_ptrs)

                inputs[f"{lora_module}_lora_ranks_{layer_idx}"] = torch.IntTensor(lora_ranks_)
                inputs[f"{lora_module}_lora_weights_pointers_{layer_idx}"] = torch.LongTensor(
                    lora_ptrs_
                )
        return inputs