Added custom mamba op and fix the mamba cache issue #1521
Conversation
| from transformers.utils import ( | ||
| ModelOutput, | ||
| logging, | ||
| ) | ||
|
|
||
| from pathlib import Path | ||
| import os | ||
| base_dir = "/workspace/custom_op_pscan_all" |
There was a problem hiding this comment.
@libinta what would be the best way to set this without hardcoding?
Atleast an env var?
There was a problem hiding this comment.
Is this dir generated on the fly? Or is it supposed to be downloaded (e.g. as part of an example) ?
There was a problem hiding this comment.
I will change based on our relative folder location
There was a problem hiding this comment.
Added HABANA_CUSTOM_OP_DIR for custom op lib folder or using the current folder as lib folder.
| A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) | ||
| ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, | ||
| and is why Mamba is called **selective** state spaces) | ||
| """ |
There was a problem hiding this comment.
@zzhang37 can you plz add a comment in the code about the different between this and original impl?
|
|
||
| # fmt: off | ||
| def slow_forward(self, input_states, cache_params: Optional[MambaCache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.LongTensor] = None): | ||
| batch_size, seq_len, _ = input_states.shape |
There was a problem hiding this comment.
@zzhang37 , can u plz add a brief code comment about the difference between this and original.
is it only Run_Mamba_Forward_Gaudi ?
|
@zzhang37 Also, all synapse dependencies merged in to 1.19 release? |
What does this PR do?
Added custom mamba pscan op and fixed the mamba cache issue
Fixes # (issue)
Before submitting