From d57671527abc5b02df750f9fa3dd9f972ce2190d Mon Sep 17 00:00:00 2001
From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Date: Thu, 30 Apr 2026 15:23:39 -0700
Subject: [PATCH] Fix LFM2 ShortConv Mamba State Indexing (#23975)

---
 python/sglang/srt/models/lfm2.py     | 9 ++++++---
 python/sglang/srt/models/lfm2_moe.py | 9 ++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

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,