Skip to content

[Mamba] eliminate D2H if tracking mamba states #20522

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ispobock merged 9 commits into from
Apr 7, 2026
Merged

[Mamba] eliminate D2H if tracking mamba states #20522

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
f59a7c1
[Mamba] eliminate D2H if tracking mamba states
Henson-Zh-Ali Mar 13, 2026
7f34349
format
Henson-Zh-Ali Mar 15, 2026
8baa202
Merge branch 'main' into opt_mamba_tracking
hzh0425 Mar 18, 2026
1f51282
pre-commit
Henson-Zh-Ali Mar 19, 2026
d9db822
Merge branch 'main' into opt_mamba_tracking
hzh0425 Mar 23, 2026
7c1384a
Merge branch 'main' into opt_mamba_tracking
hzh0425 Mar 31, 2026
3e8ed55
Merge branch 'main' into opt_mamba_tracking
hzh0425 Apr 2, 2026
fc3013a
rename conv_states_indices to conv_states_mask_indices
Henson-Zh-Ali Apr 3, 2026
c3acf52
fix lint
hzh0425 Apr 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,11 @@ def _forward_metadata(self, forward_batch: ForwardBatch):
else:
raise ValueError(f"Invalid forward mode: {forward_batch.forward_mode=}")

has_mamba_track_mask = bool(
forward_batch.mamba_track_mask is not None
and forward_batch.mamba_track_mask.any()

@yizhang2077 yizhang2077 Apr 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we also need to fix D2H here

@Henson-Zh-Ali Henson-Zh-Ali Apr 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this mean that it is pre-computed while preparing ScheduleBatch and passed to here? We think that this approach involves modifying too many files and does not yield significant performance gain. If you believe this is necessary, we will proceed with it. Thank you for your suggestion.

)
Comment thread
Henson-Zh-Ali marked this conversation as resolved.

return ForwardMetadata(
query_start_loc=query_start_loc,
mamba_cache_indices=mamba_cache_indices,
Expand All @@ -228,6 +233,7 @@ def _forward_metadata(self, forward_batch: ForwardBatch):
track_ssm_h_dst=track_ssm_h_dst,
track_ssm_final_src=track_ssm_final_src,
track_ssm_final_dst=track_ssm_final_dst,
has_mamba_track_mask=has_mamba_track_mask,
)

def init_forward_metadata(self, forward_batch: ForwardBatch):
Expand Down Expand Up @@ -613,10 +619,7 @@ def _track_mamba_state_extend(
Note: Conv state tracking for extend is handled separately via gather operations
using indices computed by `_init_track_conv_indices`.
"""
if (
forward_batch.mamba_track_mask is not None
and forward_batch.mamba_track_mask.any()
):
if forward_metadata.has_mamba_track_mask:
h = h.squeeze(0)

if forward_metadata.track_ssm_h_src.numel() > 0:
Expand Down
18 changes: 11 additions & 7 deletions python/sglang/srt/layers/attention/linear/gdn_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,15 @@ def __init__(self, model_runner: ModelRunner):
prefill_backend = get_linear_attn_prefill_backend()
self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)

def init_forward_metadata(self, forward_batch: ForwardBatch):
super().init_forward_metadata(forward_batch)
if self.forward_metadata.has_mamba_track_mask:
self.forward_metadata.mamba_track_mask_indices = \
forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
Comment thread
Henson-Zh-Ali marked this conversation as resolved.
Outdated
self.forward_metadata.conv_states_indices = forward_batch.mamba_track_indices[
self.forward_metadata.mamba_track_mask_indices
]

def forward_decode(
self,
layer: RadixLinearAttention,
Expand Down Expand Up @@ -336,16 +345,11 @@ def forward_extend(
mixed_qkv = mixed_qkv_processed.transpose(1, 2).view(seq_len, -1)
else:
mixed_qkv = mixed_qkv.transpose(0, 1)
if (
forward_batch.mamba_track_mask is not None
and forward_batch.mamba_track_mask.any()
):
conv_dst = forward_batch.mamba_track_indices
if forward_metadata.has_mamba_track_mask:
mixed_qkv_to_track = mixed_qkv[
:, forward_metadata.track_conv_indices
].transpose(0, 1)
mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track
conv_states[forward_metadata.conv_states_indices] = mixed_qkv_to_track

mixed_qkv = causal_conv1d_fn(
mixed_qkv,
Expand Down
4 changes: 4 additions & 0 deletions python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,10 @@ class ForwardMetadata:
is_target_verify: bool = False
draft_token_num: int = 1

has_mamba_track_mask: bool = False
mamba_track_mask_indices: Optional[torch.Tensor] = None
conv_states_indices: Optional[torch.Tensor] = None
Comment thread
yizhang2077 marked this conversation as resolved.
Outdated


@dataclass(kw_only=True)
class Mamba2Metadata(ForwardMetadata):
Expand Down
Loading