Repository navigation
Enable FlashInfer checkpointing SSU STP no-fallback path - #43439
danielafrimi wants to merge 6 commits into
Conversation
Signed-off-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
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 |
![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 Mamba models, introducing a configurable checkpoint_interval and expanding the Mamba KV cache to include replay and tracker tensors. The changes update the MambaMixer2 layer and FlashInferSSUBackend to utilize new FlashInfer kernels, while also extending KVBlockZeroer to support Mamba state. Critical feedback was provided regarding the efficiency of the backend implementation, specifically noting performance bottlenecks caused by manual loops, frequent memory allocations, and torch.cat operations. Furthermore, a potential race condition was identified in the slot-copying logic during CUDA stream capture when source and destination indices overlap.
| if use_grouped_checkpointing: | ||
| for group_idx in range(n_groups): | ||
| group_head_start = group_idx * heads_per_group | ||
| group_head_end = group_head_start + heads_per_group | ||
| group_slice = slice(group_idx, group_idx + 1) | ||
| old_B_snapshot = old_B[:, :, :, group_slice, :].contiguous() | ||
| old_B_group = torch.empty_like(old_B_snapshot) | ||
| for head_start in range(group_head_start, group_head_end): | ||
| head_end = head_start + 1 | ||
| head_slice = slice(head_start, head_end) | ||
| state_group = state[:, head_slice].contiguous() | ||
| old_x_group = old_x[:, :, head_slice].contiguous() | ||
| old_B_group.copy_(old_B_snapshot) | ||
| old_dt_group = old_dt[:, :, head_slice, :].contiguous() | ||
| old_cumAdt_group = old_cumAdt[ | ||
| :, :, head_slice, : | ||
| ].contiguous() | ||
| out_group = torch.empty_like(x_ckpt[:, :, head_slice, :]) | ||
| self._checkpointing_kernel( | ||
| state_group, | ||
| old_x_group, | ||
| old_B_group, | ||
| old_dt_group, | ||
| old_cumAdt_group, | ||
| cache_buf_idx, | ||
| prev_num_accepted_tokens, | ||
| x_ckpt[:, :, head_slice, :].contiguous(), | ||
| dt_ckpt[:, :, head_slice, :], | ||
| A[head_slice], | ||
| B_ckpt[:, :, group_slice, :].contiguous(), | ||
| C_ckpt[:, :, group_slice, :].contiguous(), | ||
| out_group, | ||
| D=D[head_slice] if D is not None else None, | ||
| z=( | ||
| z_ckpt[:, :, head_slice, :].contiguous() | ||
| if z_ckpt is not None | ||
| else None | ||
| ), | ||
| dt_bias=( | ||
| dt_bias[head_slice] if dt_bias is not None else None | ||
| ), | ||
| dt_softplus=dt_softplus, | ||
| state_batch_indices=kernel_state_indices, | ||
| pad_slot_id=null_block_id, | ||
| rand_seed=rand_seed, | ||
| philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds | ||
| or 10, | ||
| cu_seqlens=ckpt_cu_seqlens, | ||
| max_seqlen=kernel_max_seqlen, | ||
| ) | ||
| state[:, head_slice].copy_(state_group) | ||
| old_x[:, :, head_slice].copy_(old_x_group) | ||
| old_dt[:, :, head_slice, :].copy_(old_dt_group) | ||
| old_cumAdt[:, :, head_slice, :].copy_(old_cumAdt_group) | ||
| out_ckpt[:, :, head_slice, :].copy_(out_group) | ||
| old_B[:, :, :, group_slice, :].copy_(old_B_group) |
There was a problem hiding this comment.
The manual loop over groups and heads in _run_checkpointing_kernel is highly inefficient. It performs multiple .contiguous() calls on the state tensor (which copies the entire state for all cache slots) and allocates new tensors using torch.empty_like inside the nested loop. For a model with 128 heads, this results in 128 kernel launches and 128 sets of large memory allocations/copies per layer per step. This will significantly degrade performance and increase memory pressure. The logic should be refactored to use a single kernel launch or at least avoid per-head allocations and full-state copies.
| x_chunk = torch.cat(x_parts, dim=0) | ||
| dt_chunk = torch.cat(dt_parts, dim=0) | ||
| B_chunk = torch.cat(B_parts, dim=0) | ||
| C_chunk = torch.cat(C_parts, dim=0) | ||
| z_chunk = torch.cat(z_parts, dim=0) if z_parts is not None else None |
There was a problem hiding this comment.
Using torch.cat inside the chunking loop for varlen sequences is inefficient as it creates new tensors and involves additional memory copies on every chunk iteration. This can be a performance bottleneck for long sequences or large batches. Consider using pre-allocated buffers or refactoring the kernel to handle varlen sequences more natively without concatenation.
| _copy_checkpointing_slots_kernel[ | ||
| (n_slots, triton.cdiv(slot_size, block)) | ||
| ]( | ||
| tensor, | ||
| src_indices, | ||
| dst_indices, | ||
| slot_size, | ||
| slot_stride, | ||
| pad_slot_id, | ||
| BLOCK=block, | ||
| ) | ||
| continue |
There was a problem hiding this comment.
The _copy_checkpointing_slots_kernel is used when no scratch buffer is available, which occurs during CUDA stream capture (as seen in _get_copy_scratch at line 977). This kernel performs a direct batched copy that is unsafe when source and destination indices overlap (e.g., moving state from slot A to B while another move goes from B to C). Since Triton does not guarantee the execution order of programs, this can lead to race conditions where a slot is overwritten before its original value is read. To ensure correctness during CUDA graph capture, the scratch buffer should be pre-allocated during the warmup/profile phase so it is available during capture.
This branch has not been deployed
1 participant
Summary
This PR enables the FlashInfer checkpointing SSU path for STP without falling back to the older SSU implementation.
ssu_dispatch.py.tests/kernels/mamba/test_checkpointing_ssu_stp.py, plus worker utility coverage for the new config/routing.Validation
0.92342683850.9302501895/lustre/fsw/coreai_nvfm_llm/dafrimi/vllm-stp-no-fallback-artifacts/evals/gsm8k-safe-window-normal-cg-full-20260522_000209/nemotron-stp-safe-window-normal-cg/results_2026-05-22T00-55-08.305740.jsonFor comparison, old FlashInfer STP full runs were in the same band:
gsm8k-stp-oldfi-full-conc50-20260520: strict0.9188779378, flexible0.9310083397gsm8k-stp-oldfi-full-rerun2-conc50-nobs-20260520: strict0.9166034875, flexible0.9302501895Notes
flashinfer-checkpointing-ssu-stp-no-fallback-cleanee869a7e32a205c7ed49ca2461acc8913a1614acflashinfer-checkpointing-ssu-mtp-new-kernel.