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

LFM2: pass has_initial_state to causal_conv1d_fn for prefill (#24717)

Browse Source 
Co-authored-by: Piotr Mazurek <piotr.mazurek@liquid.ai>
This commit is contained in:
Piotr Mazurek
2026-05-13 21:57:58 -07:00
committed by GitHub
co-authored by Piotr Mazurek
parent 5c11c2492f
commit a6a6c3119b
2 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/lfm2.py
+3 -1
View File
@@ -302,11 +302,13 @@ class Lfm2ShortConv(nn.Module):
] ]
) )
cache_indices = mamba_indices.to(torch.int32) cache_indices = mamba_indices.to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens > 0
else: else:
query_start_loc = torch.tensor( query_start_loc = torch.tensor(
[0, T], dtype=torch.int32, device=hidden_states.device [0, T], dtype=torch.int32, device=hidden_states.device
) )
cache_indices = mamba_indices[:1].to(torch.int32) cache_indices = mamba_indices[:1].to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens[:1] > 0
conv_out = causal_conv1d_fn( conv_out = causal_conv1d_fn(
Bx_t, Bx_t,
@@ -314,7 +316,7 @@ class Lfm2ShortConv(nn.Module):
self.conv_bias, self.conv_bias,
query_start_loc=query_start_loc, query_start_loc=query_start_loc,
cache_indices=cache_indices, cache_indices=cache_indices,
has_initial_state=None, has_initial_state=has_initial_state,
conv_states=conv_state, conv_states=conv_state,
activation=None, activation=None,
).transpose(0, 1) ).transpose(0, 1)
python/sglang/srt/models/lfm2_moe.py
+3 -1
View File
@@ -360,10 +360,12 @@ class Lfm2MoeShortConv(nn.Module):
query_start_loc[:-1] = extend_start_loc query_start_loc[:-1] = extend_start_loc
query_start_loc[-1] = T query_start_loc[-1] = T
cache_indices = mamba_indices.to(torch.int32) cache_indices = mamba_indices.to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens > 0
else: else:
# Single sequence: [0, T] # Single sequence: [0, T]
query_start_loc = hidden_states.new_tensor([0, T], dtype=torch.int32) query_start_loc = hidden_states.new_tensor([0, T], dtype=torch.int32)
cache_indices = mamba_indices[:1].to(torch.int32) cache_indices = mamba_indices[:1].to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens[:1] > 0
conv_out = causal_conv1d_fn( conv_out = causal_conv1d_fn(
Bx_t, Bx_t,
@@ -371,7 +373,7 @@ class Lfm2MoeShortConv(nn.Module):
self.conv_bias, self.conv_bias,
query_start_loc=query_start_loc, query_start_loc=query_start_loc,
cache_indices=cache_indices, cache_indices=cache_indices,
has_initial_state=None, has_initial_state=has_initial_state,
conv_states=conv_state, conv_states=conv_state,
activation=None, activation=None,
).transpose(0, 1) ).transpose(0, 1)