diff --git a/python/sglang/srt/models/lfm2.py b/python/sglang/srt/models/lfm2.py index 7a8670aba..3694e32fa 100644 --- a/python/sglang/srt/models/lfm2.py +++ b/python/sglang/srt/models/lfm2.py @@ -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, diff --git a/python/sglang/srt/models/lfm2_moe.py b/python/sglang/srt/models/lfm2_moe.py index 2e164591e..4c7d5d06d 100644 --- a/python/sglang/srt/models/lfm2_moe.py +++ b/python/sglang/srt/models/lfm2_moe.py @@ -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,