transformers
979041cc - Register the remaining mamba-ssm kernel layers on XPU (#49035)

Commit
5 days ago
Register the remaining mamba-ssm kernel layers on XPU (#49035) Follow-up to #48702, which only mapped causal_conv1d/RMSNormGated/KDA on XPU. The mamba-ssm scan and selective-state layers were left CUDA-only, so on XPU they fell back to the reference PyTorch implementation of mamba2_chunk_scan, which materializes a large fp32 intermediate and OOMs on Bamba-9B. kernels-community/mamba-ssm v3 ships XPU builds that export these layers, so map them the same way as CUDA.
Author
Parents
Loading