[V1][Mamba1] - FP32 SSM Kernel Support (#23506)

Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
This commit is contained in:
Asaf Joseph Gardin
2025-09-02 06:53:00 +03:00
committed by GitHub
parent 0235103cbb
commit 2b41cbbf03
3 changed files with 65 additions and 32 deletions

View File

@@ -65,6 +65,11 @@ V0_UNSUPPORTED_MODELS = [
"LiquidAI/LFM2-1.2B",
]
FP32_STATE_MODELS = [
"state-spaces/mamba-130m-hf",
"Zyphra/Zamba2-1.2B-instruct",
]
# Avoid OOM
MAX_NUM_SEQS = 4
@@ -434,7 +439,7 @@ def test_full_cuda_graph(
)
@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
@pytest.mark.parametrize("model", FP32_STATE_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
def test_fp32_state(