A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Bugfix] Fix tests/kernels/test_mamba_ssm_ssd.py (#16623)
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
Tyler Michael Smith committed
Apr 15, 2025 at 01:35 UTC
dbb036cf612a3c9943254182af40597ec107be08
1 file changed
+12
-5
tests/kernels/test_mamba_ssm_ssd.py
+12
-5
@@ -5,6 +5,8 @@ import torch
5
import torch.nn.functional as F
6
from einops import rearrange, repeat
7
8
+from vllm.model_executor.layers.mamba.mamba2_metadata import (
9
+ _seq_idx_to_chunk_indices_offsets)
10
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
11
mamba_chunk_scan_combined)
12
from vllm.platforms import current_platform
@@ -160,14 +162,14 @@ def generate_continous_batched_examples(example_lens_by_batch,
162
163
# get the metadata
164
cu_seqlens = torch.tensor((0, ) + spec, device=device).cumsum(dim=0)
163
- sed_idx = torch.zeros(cu_seqlens[-1],
165
+ seq_idx = torch.zeros(cu_seqlens[-1],
166
dtype=torch.int32,
167
device=cu_seqlens.device)
168
for i, (srt, end) in enumerate(zip(
169
cu_seqlens,
170
cu_seqlens[1:],
171
)):
170
- sed_idx[srt:end] = i
172
+ seq_idx[srt:end] = i
173
174
# for cont batch
175
if IND_E is None:
@@ -177,7 +179,7 @@ def generate_continous_batched_examples(example_lens_by_batch,
179
IND_E = [end_boundary(x + y) for x, y in zip(IND_S, spec)]
180
181
yield ([Y_min[s, IND_S[s]:IND_E[s]] for s in range(num_examples)],
180
- cu_seqlens, sed_idx.unsqueeze(0), (A, dt2, X2, B2, C2))
182
+ cu_seqlens, seq_idx.unsqueeze(0), (A, dt2, X2, B2, C2))
183
184
185
@pytest.mark.parametrize("itype",
@@ -266,12 +268,15 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
268
exhausted: dict = {} # map: eg -> boolean indicating example is exhausted
269
270
states = None
269
- for Y_min, cu_seqlens, sed_idx, (A, dt, X, B,
271
+ for Y_min, cu_seqlens, seq_idx, (A, dt, X, B,
272
C) in generate_continous_batched_examples(
273
cases, num_examples, seqlen,
274
last_taken, exhausted, n_heads,
275
d_head, itype):
276
277
+ chunk_indices, chunk_offsets = _seq_idx_to_chunk_indices_offsets(
278
+ seq_idx, chunk_size)
279
+
280
Y, new_states = mamba_chunk_scan_combined(
281
X,
282
dt,
@@ -281,7 +286,9 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
286
chunk_size,
287
D=None,
288
cu_seqlens=cu_seqlens,
284
- seq_idx=sed_idx,
289
+ seq_idx=seq_idx,
290
+ chunk_indices=chunk_indices,
291
+ chunk_offsets=chunk_offsets,
292
return_varlen_states=True,
293
initial_states=states,
294
)