[Mamba1] - Kernel Level Chunk Alignment for Prefix Caching (#34798)

Signed-off-by: Josephasafg <ajgard7@gmail.com>
2026-03-01 14:40:23 +02:00
parent da543d1abe
commit bbf81f9a92
11 changed files with 251 additions and 146 deletions
--- a/tests/kernels/mamba/test_mamba_ssm.py
+++ b/tests/kernels/mamba/test_mamba_ssm.py
@@ -183,6 +183,8 @@ def selective_scan_opcheck_fn(
    block_idx_first_scheduled_token=None,
    block_idx_last_scheduled_token=None,
    initial_state_idx=None,
+    cu_chunk_seqlen=None,
+    last_chunk_indices=None,
 ):
    """if return_last_state is True, returns (out, last_state)
    last_state has shape (batch, dim, dstate).
@@ -231,6 +233,8 @@ def selective_scan_opcheck_fn(
            block_idx_first_scheduled_token,
            block_idx_last_scheduled_token,
            initial_state_idx,
+            cu_chunk_seqlen,
+            last_chunk_indices,
        ),
        test_utils=["test_schema", "test_faketensor"],
    )