minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Bugfix] Track DFlash Mamba state at checkpoint boundaries (#37818)

Browse Source 
Co-authored-by: lvweiv <vwei.lv@gmail.com>
Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
simple-sun
2026-09-11 18:38:09 -07:00
committed by GitHub
co-authored by lvweiv kpham-sgl Claude Fable 5.1 Baizhou Zhang
parent 29cc9d2dc9
commit b805cc5014
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/dflash_worker_v2.py
+6 -2
View File
@@ -1777,6 +1777,7 @@ class DFlashWorkerV2(BaseSpecWorker):
*, *,
batch: ScheduleBatch, batch: ScheduleBatch,
seq_lens_pre_verify: torch.Tensor, seq_lens_pre_verify: torch.Tensor,
seq_lens_post_verify: torch.Tensor,
commit_lens: torch.Tensor, commit_lens: torch.Tensor,
) -> None: ) -> None:
"""Commit Mamba intermediate states for accepted verify steps. """Commit Mamba intermediate states for accepted verify steps.
@@ -1796,10 +1797,10 @@ class DFlashWorkerV2(BaseSpecWorker):
mamba_track_interval = mamba_track_grid(batch.tree_cache.page_size) mamba_track_interval = mamba_track_grid(batch.tree_cache.page_size)
to_track_mask = ( to_track_mask = (
seq_lens_pre_verify // mamba_track_interval seq_lens_pre_verify // mamba_track_interval
!= batch.seq_lens // mamba_track_interval != seq_lens_post_verify // mamba_track_interval
) )
tracking_point = ( tracking_point = (
batch.seq_lens // mamba_track_interval * mamba_track_interval seq_lens_post_verify // mamba_track_interval * mamba_track_interval
) )
to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0) to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0)
can_track_mask = to_track_mask & ( can_track_mask = to_track_mask & (
@@ -2460,9 +2461,12 @@ class DFlashWorkerV2(BaseSpecWorker):
if self._need_mamba_verify_commit: if self._need_mamba_verify_commit:
assert seq_lens_pre_verify is not None assert seq_lens_pre_verify is not None
if new_seq_lens is None:
new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype)
self._update_target_mamba_state_after_verify( self._update_target_mamba_state_after_verify(
batch=batch, batch=batch,
seq_lens_pre_verify=seq_lens_pre_verify, seq_lens_pre_verify=seq_lens_pre_verify,
seq_lens_post_verify=new_seq_lens,
commit_lens=commit_lens, commit_lens=commit_lens,
) )