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

Fix LFM2 ShortConv Mamba State Indexing (#23975)

Browse Source
This commit is contained in:
Hubert Lu
2026-04-30 15:23:39 -07:00
committed by GitHub
parent 2e027b1afe
commit d57671527a
2 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/lfm2.py
+6 -3
View File
@@ -266,6 +266,9 @@ class Lfm2ShortConv(nn.Module):
layer_cache = forward_batch.req_to_token_pool.mamba2_layer_cache(self.layer_idx)
conv_state = layer_cache.conv[0]
req_pool_indices = forward_batch.req_pool_indices
mamba_indices = forward_batch.req_to_token_pool.get_mamba_indices(
req_pool_indices
)
# Project and split into gates: B (pre-conv), C (post-conv), x (input)
proj, _ = self.in_proj(hidden_states)
@@ -280,7 +283,7 @@ class Lfm2ShortConv(nn.Module):
self.conv_weight,
self.conv_bias,
activation=None,
conv_state_indices=req_pool_indices.to(torch.int32),
conv_state_indices=mamba_indices.to(torch.int32),
)
else:
# Prefill: multiple tokens, use varlen kernel
@@ -298,12 +301,12 @@ class Lfm2ShortConv(nn.Module):
),
]
)
cache_indices = req_pool_indices.to(torch.int32)
cache_indices = mamba_indices.to(torch.int32)
else:
query_start_loc = torch.tensor(
[0, T], dtype=torch.int32, device=hidden_states.device
)
cache_indices = req_pool_indices[:1].to(torch.int32)
cache_indices = mamba_indices[:1].to(torch.int32)
conv_out = causal_conv1d_fn(
Bx_t,
python/sglang/srt/models/lfm2_moe.py
+6 -3
View File
@@ -329,6 +329,9 @@ class Lfm2MoeShortConv(nn.Module):
layer_cache = forward_batch.req_to_token_pool.mamba2_layer_cache(self.layer_idx)
conv_state = layer_cache.conv[0]
req_pool_indices = forward_batch.req_pool_indices
mamba_indices = forward_batch.req_to_token_pool.get_mamba_indices(
req_pool_indices
)
proj, _ = self.in_proj(hidden_states)
B_gate, C_gate, x = proj.chunk(3, dim=-1)
@@ -341,7 +344,7 @@ class Lfm2MoeShortConv(nn.Module):
self.conv_weight,
self.conv_bias,
activation=None,
conv_state_indices=req_pool_indices.to(torch.int32),
conv_state_indices=mamba_indices.to(torch.int32),
)
else:
T = hidden_states.shape[0]
@@ -356,11 +359,11 @@ class Lfm2MoeShortConv(nn.Module):
query_start_loc = extend_start_loc.new_empty(len(extend_start_loc) + 1)
query_start_loc[:-1] = extend_start_loc
query_start_loc[-1] = T
cache_indices = req_pool_indices.to(torch.int32)
cache_indices = mamba_indices.to(torch.int32)
else:
# Single sequence: [0, T]
query_start_loc = hidden_states.new_tensor([0, T], dtype=torch.int32)
cache_indices = req_pool_indices[:1].to(torch.int32)
cache_indices = mamba_indices[:1].to(torch.int32)
conv_out = causal_conv1d_fn(
Bx_t,