from trt_execution import TRTModel import tensorrt as trt import torch from threading import Lock class MusicalPositionEmbedTransformerAccelerated: def __init__(self, params, trt_path, trt_one_step_path, encoder): self.params = params self.inference_batch_size = params.inference_batch_size self.logger = trt.Logger(trt.Logger.INFO) self.runtime = trt.Runtime(self.logger) self.stream = torch.cuda.Stream() self.model = TRTModel(self.runtime, trt_path) self.context = self.model.create_execution_context(stream=self.stream) self.model_one_step = TRTModel(self.runtime, trt_one_step_path) self.context_one_step = self.model_one_step.create_execution_context( stream=self.stream ) self.encoder = encoder self.encoder_lock = Lock() self.cache_growth_factor = 1.5 self.cache_offset = 0 self.hint_inference_batch_size(self.inference_batch_size) def get_device(self): return self.context.torch_device def forward(self, *args, **kwargs): with torch.cuda.stream(self.stream): return self.forward_impl(*args, **kwargs) def hint_inference_batch_size(self, batch_size): idx = self.model_one_step.get_profile_index_for_dim_constraint( "x", 0, batch_size ) self.context_one_step.set_optimization_profile_index(idx) self.inference_batch_size = self.model_one_step.get_profile_max_size( idx, "x", 0 ) self.a_xks = None self.a_xvs = None self.c_xks = None self.c_xvs = None self.encoder_valid = None self.symbols_consumed = 0 self.x_upload = torch.zeros( self.inference_batch_size, 1, 9, device="cpu", dtype=torch.int32, pin_memory=True, requires_grad=False, ) self.x_uploaded = torch.zeros( self.inference_batch_size, 1, 9, device=self.stream.device, dtype=torch.int32, requires_grad=False, ) def rearrange(self, idxs): assert self.symbols_consumed > 0 if self.c_xks is not None: self.c_xks[:, : idxs.shape[0]] = self.c_xks[:, idxs] self.c_xvs[:, : idxs.shape[0]] = self.c_xvs[:, idxs] self.encoder_valid[: idxs.shape[0]] = self.encoder_valid[idxs] self.a_xks[:, : idxs.shape[0]] = self.a_xks[:, idxs] self.a_xvs[:, : idxs.shape[0]] = self.a_xvs[:, idxs] def forward_impl( self, x, start_pos=0, encoder_input_ids=None, encoder_attention_mask=None, return_embeddings=False, ): assert len(x.shape) == 3 if start_pos == 0: x = x.type(torch.int32).contiguous().to(self.context.torch_device) assert encoder_input_ids is not None assert encoder_attention_mask is not None assert ( len(encoder_input_ids.shape) == 2 and encoder_input_ids.shape[0] == x.shape[0] ) encoder_input_ids = encoder_input_ids.to(self.context.torch_device) encoder_attention_mask = encoder_attention_mask.to( self.context.torch_device ) encoder_valid = encoder_attention_mask.to(dtype=torch.bool) self.encoder_valid = torch.zeros( self.inference_batch_size, encoder_valid.shape[1], device=encoder_valid.device, dtype=torch.bool, ) self.encoder_valid[: encoder_valid.shape[0]] = encoder_valid with torch.no_grad(), self.encoder_lock: encoder_out = self.encoder( encoder_input_ids, attention_mask=encoder_attention_mask ) self.stream.synchronize() self.context.set_optimization_profile_index( self.model.get_profile_index_for_dim_constraint("x", 0, x.shape[0]) ) out = self.context.eval( x=x, encoder_out=encoder_out, encoder_valid=encoder_valid, ) if return_embeddings: return out["embeddings"] self.symbols_consumed = x.shape[1] a_xks_out = out["a_xks"] a_xvs_out = out["a_xvs"] self.cache_offset = a_xks_out.shape[1] - self.symbols_consumed cache_len = self.symbols_consumed + self.cache_offset cache_size = max(int(cache_len * self.cache_growth_factor), 50) c_xks_out = out.get("c_xks", None) c_xvs_out = out.get("c_xvs", None) if c_xks_out is not None: self.c_xks = torch.zeros( c_xks_out.shape[0], self.inference_batch_size, *c_xks_out.shape[2:], device=c_xks_out.device, requires_grad=False, ) self.c_xvs = torch.zeros( c_xvs_out.shape[0], self.inference_batch_size, *c_xvs_out.shape[2:], device=c_xvs_out.device, requires_grad=False, ) self.c_xks[:, : x.shape[0]] = c_xks_out[:, : x.shape[0]] self.c_xvs[:, : x.shape[0]] = c_xvs_out[:, : x.shape[0]] self.a_xks = torch.zeros( cache_size, self.inference_batch_size, *a_xks_out.shape[2:], device=a_xks_out.device, requires_grad=False, ) self.a_xvs = torch.zeros( cache_size, self.inference_batch_size, *a_xvs_out.shape[2:], device=a_xvs_out.device, requires_grad=False, ) # one step model expects sequence first then batch self.a_xks[:cache_len, : x.shape[0]] = a_xks_out.transpose(0, 1) self.a_xvs[:cache_len, : x.shape[0]] = a_xvs_out.transpose(0, 1) return out["y"] assert self.symbols_consumed > 0 assert x.shape[0] <= self.inference_batch_size orig_batch_size = x.shape[0] if x.shape[0] < self.inference_batch_size: x = torch.cat( [ x, torch.zeros( self.inference_batch_size - x.shape[0], *x.shape[1:], dtype=x.dtype, device=x.device, ), ], dim=0, ) self.x_upload[:, :, :] = x self.x_uploaded.copy_(self.x_upload, non_blocking=True) assert start_pos <= self.symbols_consumed total_space = self.a_xks.shape[0] cache_len = start_pos + self.cache_offset if cache_len + 1 > total_space: new_total_space = int(total_space * self.cache_growth_factor) # grow the cache print(f"growing cache {total_space} -> {new_total_space} symbols") new_a_xks = torch.zeros( new_total_space, *self.a_xks.shape[1:], device=self.a_xks.device ) new_a_xvs = torch.zeros( new_total_space, *self.a_xvs.shape[1:], device=self.a_xvs.device ) new_a_xks[:total_space] = self.a_xks new_a_xvs[:total_space] = self.a_xvs self.a_xks = new_a_xks self.a_xvs = new_a_xvs self.context_one_step.eval_async( x=self.x_uploaded, a_xks=self.a_xks[:cache_len], a_xvs=self.a_xvs[:cache_len], c_xks=self.c_xks, c_xvs=self.c_xvs, encoder_valid=self.encoder_valid, ) def thunk(): with torch.cuda.stream(self.stream): out = self.context_one_step.eval_await() self.a_xks[cache_len : cache_len + 1] = out["a_xk_news"] self.a_xvs[cache_len : cache_len + 1] = out["a_xv_news"] self.symbols_consumed = start_pos + 1 return out["y"][:orig_batch_size] return thunk