Enable Mamba to use PScan Kernel

[slokesha]

What does this PR do?

Mamba model to use the PScan Kernel which will improve its throughput.
HL-SMI Version: hl-1.22.0-rc-fw-
Driver Version: 1.22.0-48ef525
Nic Driver Version: 1.22.0-48ef525

image: artifactory-kfs.habana-labs.com/docker-local/1.22.0/ubuntu22.04/habanalabs/pytorch-installer-2.7.1:1.22.0-543

G2:

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[karol-brejna-i]

Contains changes in the logic itself, as well as docs and tests.
Looks good.

[regisss]

@slokesha This doesn't work with SynapseAI v1.21 right?

[slokesha]

@slokesha This doesn't work with SynapseAI v1.21 right?

yes it fails with error:
File "/root/hs/optimum-habana/optimum/habana/transformers/models/mamba/modeling_mamba.py", line 29, in Run_Mamba_Forward_Gaudi
state_out_h = torch.ops.hpu.mamba_pscan(in_state_h, in_x_h, in_dt_h, in_A_h, in_B_h)
File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1232, in getattr
raise AttributeError(
AttributeError: '_OpNamespace' 'hpu' object has no attribute 'mamba_pscan'

[regisss]

Got it, let's wait for Synapse 1.22 to be released before merging this one in main then 🙂

[regisss]

Gaudi3 baselines are unchanged?

[karol-brejna-i]

@slokesha Are you still working on the PR? It looks like only a confirmation of baseline numbers is required and it could be merged.