Skip to content

Commit

Permalink
add retake workspace func
Browse files Browse the repository at this point in the history
  • Loading branch information
oelayan7 committed Oct 1, 2024
1 parent 026c62f commit ad67ede
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions deepspeed/ops/transformer/inference/op_binding/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,9 @@ def release_workspace(self):
def reset_cache(self):
return self.reset_cache_func() if self.reset_cache_func else None

def retake_workspace(self):
return self.retake_workspace_func() if self.retake_workspace_func else None

def allocate_workspace_fp32_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size,
external_cache, rank, max_out_tokens, min_out_tokens):
return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim,
Expand Down

0 comments on commit ad67ede

Please sign in to comment.