Repository navigation
[WIP] Enable FlashInfer checkpointing SSU for Mamba STP - #43102
danielafrimi wants to merge 3 commits into
Conversation
Signed-off-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements FlashInfer checkpointing Selective State Update (SSU) for Mamba2, introducing a checkpoint_interval configuration and integrating the checkpointing_ssu kernel into the FlashInfer backend. The implementation includes updates to the Mamba2 mixer for managing checkpointing trackers and cache, as well as new utility functions for stable state copying. Review feedback identifies an inconsistency in the parity tests where the initial state dtype was hardcoded instead of using the parameterized type, and suggests that the numerical tolerances should be tightened to ensure better precision in the final version.
| initial_state = 0.1 * torch.randn( | ||
| cache_size, | ||
| nheads, | ||
| head_dim, | ||
| dstate, | ||
| device=device, | ||
| dtype=torch.float16, | ||
| ) |
There was a problem hiding this comment.
The initial_state is created with a hardcoded dtype=torch.float16, while the test is parameterized with dtype=torch.bfloat16. This is inconsistent. To make the test more robust and general, consider using the parameterized dtype for the initial_state as well. This would test the case where state and inputs have the same bfloat16 dtype. If testing mixed precision is intended, it would be better to make it explicit, for example by parameterizing the state dtype separately.
| initial_state = 0.1 * torch.randn( | |
| cache_size, | |
| nheads, | |
| head_dim, | |
| dstate, | |
| device=device, | |
| dtype=torch.float16, | |
| ) | |
| initial_state = 0.1 * torch.randn( | |
| cache_size, | |
| nheads, | |
| head_dim, | |
| dstate, | |
| device=device, | |
| dtype=dtype, | |
| ) |
| torch.testing.assert_close( | ||
| new_out.float(), | ||
| old_out.float(), | ||
| atol=3e-2, | ||
| rtol=3e-2, | ||
| msg=f"STP output mismatch at decode step {step}", | ||
| ) |
There was a problem hiding this comment.
The tolerances atol=3e-2 and rtol=3e-2 are very high for a numerical parity test. This may hide regressions or indicate significant numerical differences between the two kernels. While I understand from the PR description that this is a work-in-progress and there are known mismatches, it's important to aim for much tighter tolerances (e.g., 1e-5) for the final version.
Adds STP replay coverage for non-contiguous state indices, destination slot moves, and a production-shaped reproducer for remaining checkpointing drift. Also avoids the known large-batch checkpointing_ssu drift by advancing STP replay one request at a time.\n\nValidation: python -m pytest tests/kernels/mamba/test_checkpointing_ssu_stp.py -q -> 5 passed, 1 xfailed.\n\nKnown remaining issue: live concurrent four-prompt serving with async scheduling + CUDA graph is still under investigation; full GSM8K has not been rerun on this commit yet. Signed-off-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
GSM8K limit-50 conc50 with CUDA graph + async: new checkpointing STP 0.92 strict / 0.92 flexible, matching old FI STP baseline 0.92 strict / 0.92 flexible. Signed-off-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
1 participant
Draft/WIP branch to share the current STP FlashInfer checkpointing SSU integration work.
What is included:
checkpointing_ssubuffers for Mamba2 STP decode.--mamba-checkpoint-interval 1without a separate opt-in kernel-selection flag.Current validation:
tests/v1/worker/test_mamba_utils.pypasses.tests/kernels/mamba/test_checkpointing_ssu_stp.pypassed in the cluster container.Known blocker:
Notes: