@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fea3e476aa
Fork 0

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 14, 05:30 UTC

[Kernel] Chunk-aligned mamba2 (#24683)

Thomas Parnell committed Sep 29, 2025 at 23:18 UTC fea3e476aa3eec13c39c4cb61eab88e8eb70ad7b
8 files changed +250 -434
vllm/model_executor/layers/mamba/mamba_mixer2.py
+4 -4
@@ -502,9 +502,9 @@ class MambaMixer2(MambaBase, CustomOp):
502 prep_initial_states = attn_metadata.prep_initial_states
503 chunk_size = attn_metadata.chunk_size
504 seq_idx_p = attn_metadata.seq_idx_p
505 - chunk_indices_p = attn_metadata.chunk_indices_p
506 - chunk_offsets_p = attn_metadata.chunk_offsets_p
505 query_start_loc_p = attn_metadata.query_start_loc_p
506 + cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
507 + last_chunk_indices_p = attn_metadata.last_chunk_indices_p
508
509 # 1. Gated MLP's linear projection
510 projected_states, _ = self.in_proj(hidden_states)
@@ -634,9 +634,9 @@ class MambaMixer2(MambaBase, CustomOp):
634 z=None,
635 dt_bias=self.dt_bias,
636 seq_idx=seq_idx_p,
637 - chunk_indices=chunk_indices_p,
638 - chunk_offsets=chunk_offsets_p,
637 cu_seqlens=query_start_loc_p,
638 + cu_chunk_seqlens=cu_chunk_seqlen_p,
639 + last_chunk_indices=last_chunk_indices_p,
640 initial_states=initial_states,
641 dt_softplus=True,
642 dt_limit=(0.0, float("inf")),
vllm/model_executor/layers/mamba/ops/ssd_bmm.py
+17 -25
@@ -6,8 +6,6 @@
6
7 # ruff: noqa: E501,SIM102
8
9 -import math
10 -
9 import torch
10
11 from vllm.triton_utils import tl, triton
@@ -96,7 +94,7 @@ def _bmm_chunk_fwd_kernel(
94 a_ptr,
95 b_ptr,
96 out_ptr,
99 - seq_idx_ptr,
97 + cu_chunk_seqlens_ptr,
98 # Matrix dimensions
99 seqlen,
100 chunk_size: tl.constexpr,
@@ -112,7 +110,6 @@ def _bmm_chunk_fwd_kernel(
110 stride_out_head: tl.int64,
111 stride_outm: tl.int64,
112 stride_outn: tl.constexpr,
115 - stride_seq_idx_seqlen: tl.constexpr,
113 # Meta-parameters
114 IS_CAUSAL: tl.constexpr,
115 dot_dtype: tl.constexpr,
@@ -129,10 +126,12 @@ def _bmm_chunk_fwd_kernel(
126 if IS_CAUSAL:
127 if pid_n * BLOCK_SIZE_N >= (pid_m + 1) * BLOCK_SIZE_M:
128 return
132 - a_ptr += pid_c * chunk_size * stride_a_seqlen + pid_h * stride_a_head
133 - b_ptr += pid_c * chunk_size * stride_b_seqlen + pid_h * stride_b_head
129
135 - seq_idx_ptr += pid_c * chunk_size * stride_seq_idx_seqlen
130 + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c)
131 + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1)
132 +
133 + a_ptr += chunk_seqlen_start * stride_a_seqlen + pid_h * stride_a_head
134 + b_ptr += chunk_seqlen_start * stride_b_seqlen + pid_h * stride_b_head
135
136 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
137 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
@@ -141,7 +140,7 @@ def _bmm_chunk_fwd_kernel(
140 offs_k[None, :] * stride_ak)
141 b_ptrs = b_ptr + (offs_k[:, None] * stride_bk +
142 offs_n[None, :] * stride_b_seqlen)
144 - chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size)
143 + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start
144
145 acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
146
@@ -162,16 +161,6 @@ def _bmm_chunk_fwd_kernel(
161 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
162 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
163
165 - # Zero out the results that are not from the same request
166 - # in the varlen batch
167 - seq_idx_m = tl.load(seq_idx_ptr + offs_m * stride_seq_idx_seqlen,
168 - mask=offs_m < chunk_size_limit,
169 - other=-1)
170 - seq_idx_n = tl.load(seq_idx_ptr + offs_n * stride_seq_idx_seqlen,
171 - mask=offs_n < chunk_size_limit,
172 - other=-2)
173 - acc = tl.where(seq_idx_m[:, None] == seq_idx_n[None, :], acc, 0.0)
174 -
164 out = acc.to(out_ptr.dtype.element_ty)
165 out_ptr += pid_c * stride_out_chunk + pid_h * stride_out_head
166 out_ptrs = out_ptr + (stride_outm * offs_m[:, None] +
@@ -182,12 +171,18 @@ def _bmm_chunk_fwd_kernel(
171 (offs_n[None, :] < chunk_size))
172
173
185 -def _bmm_chunk_fwd(a, b, chunk_size, seq_idx, causal=False, output_dtype=None):
174 +def _bmm_chunk_fwd(a,
175 + b,
176 + chunk_size,
177 + cu_chunk_seqlens,
178 + causal=False,
179 + output_dtype=None):
180 """
181 Argument:
182 a: (seqlen, ngroups, k)
183 b: (seqlen, ngroups, k)
190 - seq_idx: (seqlen,). out[i, j] for seq_idx[i] != seq_idx[j] will be zeroed out.
184 + chunk_size: int
185 + cu_chunk_seq_lens: (nchunks+1,)
186 causal: if True, then out[i, j] for i > j will be arbitrary, only out[i, j] for i <= j are
187 guaranteed to be correct.
188 Return:
@@ -195,14 +190,12 @@ def _bmm_chunk_fwd(a, b, chunk_size, seq_idx, causal=False, output_dtype=None):
190 """
191 seqlen, ngroups, k = a.shape
192 assert b.shape == a.shape
198 - assert seq_idx is not None
199 - assert seq_idx.shape == (seqlen, )
193 if a.stride(-1) != 1 and a.stride(0) != 1:
194 a = a.contiguous()
195 if b.stride(-1) != 1 and b.stride(0) != 1:
196 b = b.contiguous()
197
205 - nchunks = math.ceil(seqlen / chunk_size)
198 + nchunks = len(cu_chunk_seqlens) - 1
199 # Allocates output.
200 out_dtype = a.dtype if output_dtype is None else output_dtype
201 out = torch.empty((nchunks, ngroups, chunk_size, chunk_size),
@@ -220,7 +213,7 @@ def _bmm_chunk_fwd(a, b, chunk_size, seq_idx, causal=False, output_dtype=None):
213 a_ptr=a,
214 b_ptr=b,
215 out_ptr=out,
223 - seq_idx_ptr=seq_idx,
216 + cu_chunk_seqlens_ptr=cu_chunk_seqlens,
217 seqlen=seqlen,
218 chunk_size=chunk_size,
219 K=k,
@@ -235,7 +228,6 @@ def _bmm_chunk_fwd(a, b, chunk_size, seq_idx, causal=False, output_dtype=None):
228 stride_out_head=out.stride(1),
229 stride_outm=out.stride(-2),
230 stride_outn=out.stride(-1),
238 - stride_seq_idx_seqlen=seq_idx.stride(0),
231 IS_CAUSAL=causal,
232 dot_dtype=dot_dtype,
233 )
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+79 -157
@@ -120,9 +120,7 @@ def _chunk_scan_fwd_kernel(
120 states_ptr,
121 D_ptr,
122 initstates_ptr,
123 - chunk_indices_ptr,
124 - chunk_offsets_ptr,
125 - chunk_meta_num,
123 + cu_chunk_seqlens_ptr,
124 # Matrix dimensions
125 chunk_size: tl.constexpr,
126 hdim: tl.constexpr,
@@ -149,7 +147,7 @@ def _chunk_scan_fwd_kernel(
147 stride_dA_cs_chunk: tl.int64,
148 stride_dA_cs_head: tl.int64,
149 stride_dA_cs_csize: tl.constexpr,
152 - stride_seq_idx_seqlen: tl.constexpr,
150 + stride_seq_idx_chunk: tl.constexpr,
151 stride_C_seqlen: tl.int64,
152 stride_C_head: tl.int64,
153 stride_C_dstate: tl.constexpr,
@@ -175,170 +173,107 @@ def _chunk_scan_fwd_kernel(
173 HAS_INITSTATES: tl.constexpr,
174 ):
175 pid_c = tl.program_id(axis=1).to(tl.int64)
178 - if not HAS_INITSTATES:
179 - c_idx = pid_c
180 - c_off = 0
181 - else:
182 - c_idx = tl.load(chunk_indices_ptr + pid_c, mask=pid_c > -1, other=0)
183 - c_off = tl.load(chunk_offsets_ptr + pid_c, mask=pid_c > -1, other=0)
184 -
176 pid_h = tl.program_id(axis=2)
177 num_pid_n = tl.cdiv(hdim, BLOCK_SIZE_N)
178 pid_m = tl.program_id(axis=0) // num_pid_n
179 pid_n = tl.program_id(axis=0) % num_pid_n
189 - cb_ptr += c_idx * stride_cb_chunk + (pid_h //
180 + cb_ptr += pid_c * stride_cb_chunk + (pid_h //
181 nheads_ngroups_ratio) * stride_cb_head
191 - x_ptr += c_idx * chunk_size * stride_x_seqlen + pid_h * stride_x_head
192 - dt_ptr += c_idx * stride_dt_chunk + pid_h * stride_dt_head
193 - dA_cumsum_ptr += c_idx * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
194 - C_ptr += c_idx * chunk_size * stride_C_seqlen + (
182 + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c)
183 + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1)
184 + x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head
185 + dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head
186 + dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
187 + C_ptr += chunk_seqlen_start * stride_C_seqlen + (
188 pid_h // nheads_ngroups_ratio) * stride_C_head
189
190 # M-block offsets and prev states
191 # - logic in next block may override these if there is an active offset
199 - offs_m = pid_m * BLOCK_SIZE_M + c_off + tl.arange(0, BLOCK_SIZE_M)
200 - prev_states_ptr = states_ptr + c_idx * stride_states_chunk + pid_h * stride_states_head
201 - prev_states_hdim = stride_states_hdim
202 - prev_states_dstate = stride_states_dstate
203 -
204 - chunk_size_limit = min(chunk_size, seqlen - c_idx * chunk_size)
205 -
206 - seq_idx_ptr += c_idx * chunk_size * stride_seq_idx_seqlen
207 - # - we only need seq_idx_prev to be aligned to chunk boundary
208 - seq_idx_prev = tl.load(seq_idx_ptr - stride_seq_idx_seqlen,
209 - mask=c_idx >= 1,
210 - other=0)
211 -
212 - if HAS_INITSTATES:
213 - # if there are init states, we only need seq_idx_m to point
214 - # what is the current seq_idx
215 -
216 - # get current seq idx
217 - if (pid_m * BLOCK_SIZE_M + c_off) < chunk_size_limit:
218 - seq_idx_m = tl.load(
219 - seq_idx_ptr +
220 - (pid_m * BLOCK_SIZE_M + c_off) * stride_seq_idx_seqlen, )
221 -
222 - # - recall that in ssd_state_passing, for the case c_off == 0
223 - # i.e., the very first sequence, we made states_ptr hold its initial state
224 - # so this edge case is taken care of
225 - if ((c_off == 0) and (seq_idx_prev != seq_idx_m
226 - ) # if a seq is changed exactly on boundary
227 - or (c_off > 0) # implies a new example (pseudo chunk)
228 - ):
192 + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
193 +
194 + seq_idx_ptr += pid_c * stride_seq_idx_chunk
195 + seq_idx = tl.load(seq_idx_ptr)
196 + seq_idx_prev = tl.load(seq_idx_ptr - stride_seq_idx_chunk,
197 + mask=pid_c >= 1,
198 + other=-1)
199 +
200 + if HAS_INITSTATES and (seq_idx != seq_idx_prev):
201 + prev_states_ptr = initstates_ptr + seq_idx * stride_init_states_batch + pid_h * stride_init_states_head
202 + prev_states_hdim = stride_init_states_hdim
203 + prev_states_dstate = stride_init_states_dstate
204 + else:
205 + prev_states_ptr = states_ptr + (
206 + pid_c - 1) * stride_states_chunk + pid_h * stride_states_head
207 + prev_states_hdim = stride_states_hdim
208 + prev_states_dstate = stride_states_dstate
209
230 - # - replace prev_states_ptr with init_states
231 - prev_states_ptr = initstates_ptr + seq_idx_m * stride_init_states_batch + pid_h * stride_init_states_head
232 - prev_states_hdim = stride_init_states_hdim # override strides
233 - prev_states_dstate = stride_init_states_dstate
210 + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start
211
212 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
213 dA_cs_m = tl.load(dA_cumsum_ptr + offs_m * stride_dA_cs_csize,
214 mask=offs_m < chunk_size,
215 other=0.0).to(tl.float32)
216
240 - # - handle chunk state limit
241 - if HAS_INITSTATES:
242 - # have to split this if otherwise compilation will have problems
243 - dA_cs_m_boundary = 0.0
244 -
245 - # get the c_idx for the next (logica) chunk
246 - c_idx_n = tl.load(
247 - chunk_indices_ptr + (pid_c + 1),
248 - mask=pid_c > -1 and (pid_c + 1) < chunk_meta_num,
249 - other=-1 # to trigger different chunk
250 - )
251 -
252 - # - there are things to consider
253 - # A. if c_off > 0 then we need to move the dA_cs boundary to ensure correct
254 - # contribution of past states
255 - # B. if c_off_n < chunk_size_limit, then we need to adjust this so as not to
256 - # encroach into the next sequence, where c_off_n is the offset of the next
257 - # (logical) chunk.
258 - # An equivalent check for B is c_idx == c_idx_n, where there is repetition in
259 - # (logical) chunk indices.
260 -
261 - if (c_idx == c_idx_n) or c_off > 0:
262 -
263 - # get the next offset
264 - c_off_n = tl.load(chunk_offsets_ptr + (pid_c + 1),
265 - mask=pid_c > -1 and (pid_c + 1) < chunk_meta_num,
266 - other=chunk_size)
267 -
268 - # in this case, adjust down the chunk_size_limit
269 - if c_idx == c_idx_n:
270 - chunk_size_limit = min(c_off_n, chunk_size_limit)
271 -
272 - # get the cs at the offset boundary
273 - # - c_off == 0 is a passthrough
274 - # - We need dA_cs at the boundary, defined by c_off - no need
275 - # to increase pointer by pid_m (it is a constant offset,
276 - # i.e. the same for all blocks)
277 - dA_cs_m_boundary = tl.load(
278 - dA_cumsum_ptr + (c_off - 1) * stride_dA_cs_csize,
279 - mask=(((c_off - 1) > -1) and ((c_off) < chunk_size)),
280 - other=0.0).to(tl.float32)
281 - else:
282 - # - handle seq idx when HAS_INITSTATES==False
283 - seq_idx_m = tl.load(seq_idx_ptr + offs_m * stride_seq_idx_seqlen,
284 - mask=offs_m < chunk_size_limit,
285 - other=-1)
286 -
217 acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
218
289 - # Without the if (pid_c > -1), with Triton 2.1.0, I get
290 - # Assertion `!(srcMmaLayout && dstMmaLayout) && "Unexpected mma -> mm a layout conversion"' failed.
291 - # With Triton 2.2.0, this works
292 - if IS_TRITON_22 or c_idx > -1:
293 - # Faster to just do 1 iteration with larger BLOCK_SIZE_K, up to block size 128
294 - offs_k_dstate = tl.arange(
295 - 0, BLOCK_SIZE_DSTATE if BLOCK_SIZE_DSTATE <= 128 else BLOCK_SIZE_K)
296 - C_ptrs = C_ptr + (offs_m[:, None] * stride_C_seqlen +
297 - offs_k_dstate[None, :] * stride_C_dstate)
219 + offs_out_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
220 + offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
221
299 - prev_states_ptrs = prev_states_ptr + (
300 - offs_n[None, :] * prev_states_hdim +
301 - offs_k_dstate[:, None] * prev_states_dstate)
222 + # Faster to just do 1 iteration with larger BLOCK_SIZE_K, up to block size 128
223 + offs_k_dstate = tl.arange(
224 + 0, BLOCK_SIZE_DSTATE if BLOCK_SIZE_DSTATE <= 128 else BLOCK_SIZE_K)
225 + C_ptrs = C_ptr + (offs_m[:, None] * stride_C_seqlen +
226 + offs_k_dstate[None, :] * stride_C_dstate)
227 +
228 + scale_m = tl.exp(dA_cs_m)
229 + if BLOCK_SIZE_DSTATE <= 128:
230 + C = tl.load(C_ptrs,
231 + mask=(offs_m[:, None] < chunk_size_limit) &
232 + (offs_k_dstate[None, :] < dstate),
233 + other=0.0)
234
303 - if not HAS_INITSTATES:
304 - # - this is for continuous batching where there is no init states
305 - scale_m = tl.where(seq_idx_m == seq_idx_prev, tl.exp(dA_cs_m), 0.0)
235 + if not HAS_INITSTATES and (seq_idx != seq_idx_prev):
236 + # if no init states AND starting a new sequence, we need zeros
237 + prev_states = tl.zeros((BLOCK_SIZE_DSTATE, BLOCK_SIZE_N),
238 + dtype=C_ptr.dtype.element_ty)
239 else:
307 - # - if there is initstates, we will rely on prev_states, no zeroing
308 - # required.
309 - scale_m = tl.exp(dA_cs_m - dA_cs_m_boundary)
310 -
311 - if BLOCK_SIZE_DSTATE <= 128:
312 - C = tl.load(C_ptrs,
313 - mask=(offs_m[:, None] < chunk_size_limit) &
314 - (offs_k_dstate[None, :] < dstate),
315 - other=0.0)
316 -
240 + # otherwise read the previous state
241 + prev_states_ptrs = prev_states_ptr \
242 + + offs_n[None, :] * prev_states_hdim \
243 + + offs_k_dstate[:, None] * prev_states_dstate
244 prev_states = tl.load(prev_states_ptrs,
245 mask=(offs_k_dstate[:, None] < dstate) &
246 (offs_n[None, :] < hdim),
247 other=0.0)
248 prev_states = prev_states.to(C_ptr.dtype.element_ty)
322 - acc = tl.dot(C, prev_states) * scale_m[:, None]
323 - else:
324 - for k in range(0, dstate, BLOCK_SIZE_K):
325 - C = tl.load(C_ptrs,
326 - mask=(offs_m[:, None] < chunk_size_limit) &
327 - (offs_k_dstate[None, :] < dstate - k),
328 - other=0.0)
329 - # C = (C * scale_m[:, None]).to(C_ptr.dtype.element_ty)
249 +
250 + acc = tl.dot(C, prev_states) * scale_m[:, None]
251 +
252 + else:
253 + prev_states_ptrs = prev_states_ptr \
254 + + offs_n[None, :] * prev_states_hdim \
255 + + offs_k_dstate[:, None] * prev_states_dstate
256 + for k in range(0, dstate, BLOCK_SIZE_K):
257 + C = tl.load(C_ptrs,
258 + mask=(offs_m[:, None] < chunk_size_limit) &
259 + (offs_k_dstate[None, :] < dstate - k),
260 + other=0.0)
261 + if not HAS_INITSTATES and (seq_idx != seq_idx_prev):
262 + prev_states = tl.zeros((BLOCK_SIZE_DSTATE, BLOCK_SIZE_K),
263 + dtype=C_ptr.dtype.element_ty)
264 + else:
265 prev_states = tl.load(
266 prev_states_ptrs,
267 mask=(offs_k_dstate[:, None] < dstate - k) &
268 (offs_n[None, :] < hdim),
269 other=0.0)
270 prev_states = prev_states.to(C_ptr.dtype.element_ty)
336 - acc += tl.dot(C, prev_states)
337 - C_ptrs += BLOCK_SIZE_K
338 - prev_states_ptrs += BLOCK_SIZE_K
339 - acc *= scale_m[:, None]
271 + acc += tl.dot(C, prev_states)
272 + C_ptrs += BLOCK_SIZE_K
273 + prev_states_ptrs += BLOCK_SIZE_K
274 + acc *= scale_m[:, None]
275
341 - offs_k = tl.arange(0, BLOCK_SIZE_K) + c_off
276 + offs_k = tl.arange(0, BLOCK_SIZE_K)
277 cb_ptrs = cb_ptr + (offs_m[:, None] * stride_cb_csize_m +
278 offs_k[None, :] * stride_cb_csize_k)
279 x_ptrs = x_ptr + (offs_k[:, None] * stride_x_seqlen +
@@ -375,7 +310,7 @@ def _chunk_scan_fwd_kernel(
310 dt_ptrs += BLOCK_SIZE_K * stride_dt_csize
311 dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize
312
378 - offs_out_m = pid_m * BLOCK_SIZE_M + c_off + tl.arange(0, BLOCK_SIZE_M)
313 + offs_out_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
314 offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
315
316 if HAS_D:
@@ -393,7 +328,7 @@ def _chunk_scan_fwd_kernel(
328 acc += x_residual * D
329
330 if HAS_Z:
396 - z_ptr += c_idx * chunk_size * stride_z_seqlen + pid_h * stride_z_head
331 + z_ptr += chunk_seqlen_start * stride_z_seqlen + pid_h * stride_z_head
332 z_ptrs = z_ptr + (stride_z_seqlen * offs_out_m[:, None] +
333 stride_z_hdim * offs_out_n[None, :])
334 z = tl.load(z_ptrs,
@@ -402,7 +337,7 @@ def _chunk_scan_fwd_kernel(
337 other=0.0).to(tl.float32)
338 acc *= z * tl.sigmoid(z)
339
405 - out_ptr += c_idx * chunk_size * stride_out_seqlen + pid_h * stride_out_head
340 + out_ptr += chunk_seqlen_start * stride_out_seqlen + pid_h * stride_out_head
341 out_ptrs = out_ptr + (stride_out_seqlen * offs_out_m[:, None] +
342 offs_out_n[None, :] * stride_out_hdim)
343 tl.store(out_ptrs,
@@ -418,12 +353,11 @@ def _chunk_scan_fwd(
353 dA_cumsum,
354 C,
355 states,
356 + cu_chunk_seqlens,
357 out,
358 seq_idx,
359 D=None,
360 z=None,
425 - chunk_indices=None,
426 - chunk_offsets=None,
361 initial_states=None,
362 ):
363 assert seq_idx is not None, "this implementation requires seq_idx"
@@ -441,20 +375,10 @@ def _chunk_scan_fwd(
375 assert dt.shape == (nheads, nchunks, chunk_size)
376 assert dA_cumsum.shape == (nheads, nchunks, chunk_size)
377 assert states.shape == (nchunks, nheads, headdim, dstate)
444 - assert seq_idx.shape == (seqlen, )
445 -
446 - if initial_states is not None:
447 - # with initial states, we need to take care of how
448 - # seq_idx crosses the boundaries
449 - assert chunk_indices is not None and chunk_offsets is not None, \
450 - "chunk_indices and chunk_offsets should have been set"
451 - else:
452 - chunk_indices, chunk_offsets = None, None
378 + assert seq_idx.shape == (nchunks, )
379
454 - grid = lambda META: (
455 - triton.cdiv(chunk_size, META['BLOCK_SIZE_M']) * triton.cdiv(
456 - headdim, META['BLOCK_SIZE_N']), nchunks
457 - if chunk_offsets is None else len(chunk_offsets), nheads)
380 + grid = lambda META: (triton.cdiv(chunk_size, META['BLOCK_SIZE_M']) * triton
381 + .cdiv(headdim, META['BLOCK_SIZE_N']), nchunks, nheads)
382
383 z_strides = ((z.stride(0), z.stride(1), z.stride(2)) if z is not None else
384 (0, 0, 0))
@@ -476,9 +400,7 @@ def _chunk_scan_fwd(
400 states_ptr=states,
401 D_ptr=D,
402 initstates_ptr=initial_states,
479 - chunk_indices_ptr=chunk_indices,
480 - chunk_offsets_ptr=chunk_offsets,
481 - chunk_meta_num=len(chunk_indices) if chunk_indices is not None else 0,
403 + cu_chunk_seqlens_ptr=cu_chunk_seqlens,
404 chunk_size=chunk_size,
405 hdim=headdim,
406 dstate=dstate,
@@ -503,7 +425,7 @@ def _chunk_scan_fwd(
425 stride_dA_cs_chunk=dA_cumsum.stride(1),
426 stride_dA_cs_head=dA_cumsum.stride(0),
427 stride_dA_cs_csize=dA_cumsum.stride(2),
506 - stride_seq_idx_seqlen=seq_idx.stride(0),
428 + stride_seq_idx_chunk=seq_idx.stride(0),
429 stride_C_seqlen=C.stride(0),
430 stride_C_head=C.stride(1),
431 stride_C_dstate=C.stride(2),
vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
+19 -29
@@ -6,8 +6,6 @@
6
7 # ruff: noqa: E501
8
9 -import math
10 -
9 import torch
10
11 from vllm.triton_utils import tl, triton
@@ -34,6 +32,7 @@ def _chunk_cumsum_fwd_kernel(
32 dt_bias_ptr,
33 dt_out_ptr,
34 dA_cumsum_ptr,
35 + cu_chunk_seqlens_ptr,
36 # Matrix dimension
37 seqlen,
38 nheads: tl.constexpr,
@@ -61,7 +60,11 @@ def _chunk_cumsum_fwd_kernel(
60 # https://github.com/triton-lang/triton/issues/1058
61 pid_c = tl.program_id(axis=0).to(tl.int64)
62 pid_h = tl.program_id(axis=1)
64 - dt_ptr += pid_c * chunk_size * stride_dt_seqlen
63 +
64 + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c)
65 + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1)
66 +
67 + dt_ptr += chunk_seqlen_start * stride_dt_seqlen
68 dt_out_ptr += pid_c * stride_dt_out_chunk
69 dA_cumsum_ptr += pid_c * stride_dA_cs_chunk
70
@@ -74,7 +77,7 @@ def _chunk_cumsum_fwd_kernel(
77 offs_c[None, :] * stride_dt_out_csize)
78 dA_cs_ptrs = dA_cumsum_ptr + (offs_h[:, None] * stride_dA_cs_head +
79 offs_c[None, :] * stride_dA_cs_csize)
77 - chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size)
80 + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start
81
82 dt = tl.load(dt_ptrs,
83 mask=(offs_h[:, None] < nheads) &
@@ -188,7 +191,7 @@ def _chunk_state_fwd_kernel(
191 states_ptr,
192 dt_ptr,
193 dA_cumsum_ptr,
191 - seq_idx_ptr,
194 + cu_chunk_seqlens_ptr,
195 # Matrix dimensions
196 hdim: tl.constexpr,
197 dstate: tl.constexpr,
@@ -212,7 +215,6 @@ def _chunk_state_fwd_kernel(
215 stride_dA_cs_head: tl.int64,
216 stride_dA_cs_chunk: tl.int64,
217 stride_dA_cs_csize: tl.constexpr,
215 - stride_seq_idx_seqlen: tl.constexpr,
218 # Meta-parameters
219 BLOCK_SIZE_M: tl.constexpr,
220 BLOCK_SIZE_N: tl.constexpr,
@@ -223,14 +225,14 @@ def _chunk_state_fwd_kernel(
225 num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N)
226 pid_m = tl.program_id(axis=0) // num_pid_n
227 pid_n = tl.program_id(axis=0) % num_pid_n
226 - b_ptr += pid_c * chunk_size * stride_b_seqlen + (
228 + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c)
229 + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1)
230 + b_ptr += chunk_seqlen_start * stride_b_seqlen + (
231 pid_h // nheads_ngroups_ratio) * stride_b_head
228 - x_ptr += pid_c * chunk_size * stride_x_seqlen + pid_h * stride_x_head
232 + x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head
233 dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head
234 dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
235
232 - seq_idx_ptr += pid_c * chunk_size * stride_seq_idx_seqlen
233 -
236 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
237 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
238 offs_k = tl.arange(0, BLOCK_SIZE_K)
@@ -243,10 +245,7 @@ def _chunk_state_fwd_kernel(
245 (chunk_size - 1) * stride_dA_cs_csize).to(tl.float32)
246 dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize
247
246 - seq_idx_ptrs = seq_idx_ptr + offs_k * stride_seq_idx_seqlen
247 - chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size)
248 - seq_idx_last = tl.load(seq_idx_ptr +
249 - (chunk_size_limit - 1) * stride_seq_idx_seqlen)
248 + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start
249
250 acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
251 for k in range(0, chunk_size_limit, BLOCK_SIZE_K):
@@ -261,15 +260,9 @@ def _chunk_state_fwd_kernel(
260 dA_cs_k = tl.load(dA_cumsum_ptrs,
261 mask=offs_k < chunk_size_limit - k,
262 other=0.0).to(tl.float32)
264 -
265 - seq_idx_k = tl.load(seq_idx_ptrs,
266 - mask=offs_k < chunk_size_limit - k,
267 - other=-1)
263 dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k,
264 other=0.0).to(tl.float32)
270 -
271 - scale = tl.where(seq_idx_k == seq_idx_last,
272 - tl.exp(dA_cs_last - dA_cs_k) * dt_k, 0.0)
265 + scale = tl.exp(dA_cs_last - dA_cs_k) * dt_k
266 b *= scale[:, None]
267 b = b.to(x_ptr.dtype.element_ty)
268 acc += tl.dot(x, b)
@@ -278,7 +271,6 @@ def _chunk_state_fwd_kernel(
271 b_ptrs += BLOCK_SIZE_K * stride_b_seqlen
272 dt_ptrs += BLOCK_SIZE_K * stride_dt_csize
273 dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize
281 - seq_idx_ptrs += BLOCK_SIZE_K * stride_seq_idx_seqlen
274
275 states = acc.to(states_ptr.dtype.element_ty)
276
@@ -534,6 +526,7 @@ def _chunk_state_varlen_kernel(
526 def _chunk_cumsum_fwd(dt,
527 A,
528 chunk_size,
529 + cu_chunk_seqlens,
530 dt_bias=None,
531 dt_softplus=False,
532 dt_limit=(0.0, float("inf"))):
@@ -541,7 +534,7 @@ def _chunk_cumsum_fwd(dt,
534 assert A.shape == (nheads, )
535 if dt_bias is not None:
536 assert dt_bias.shape == (nheads, )
544 - nchunks = math.ceil(seqlen / chunk_size)
537 + nchunks = cu_chunk_seqlens.shape[0] - 1
538 dt_out = torch.empty(nheads,
539 nchunks,
540 chunk_size,
@@ -561,6 +554,7 @@ def _chunk_cumsum_fwd(dt,
554 dt_bias_ptr=dt_bias,
555 dt_out_ptr=dt_out,
556 dA_cumsum_ptr=dA_cumsum,
557 + cu_chunk_seqlens_ptr=cu_chunk_seqlens,
558 seqlen=seqlen,
559 nheads=nheads,
560 chunk_size=chunk_size,
@@ -588,7 +582,7 @@ def _chunk_state_fwd(B,
582 x,
583 dt,
584 dA_cumsum,
591 - seq_idx=None,
585 + cu_chunk_seqlens,
586 states=None,
587 states_in_fp32=True):
588 seqlen, nheads, headdim = x.shape
@@ -599,9 +593,6 @@ def _chunk_state_fwd(B,
593 assert dt.shape == (nheads, nchunks, chunk_size)
594 assert dA_cumsum.shape == dt.shape
595
602 - assert seq_idx is not None
603 - assert seq_idx.shape == (seqlen, )
604 -
596 if states is not None:
597 assert states.shape == (nchunks, nheads, headdim, dstate)
598 else:
@@ -619,7 +610,7 @@ def _chunk_state_fwd(B,
610 states_ptr=states,
611 dt_ptr=dt,
612 dA_cumsum_ptr=dA_cumsum,
622 - seq_idx_ptr=seq_idx,
613 + cu_chunk_seqlens_ptr=cu_chunk_seqlens,
614 hdim=headdim,
615 dstate=dstate,
616 chunk_size=chunk_size,
@@ -641,7 +632,6 @@ def _chunk_state_fwd(B,
632 stride_dA_cs_head=dA_cumsum.stride(0),
633 stride_dA_cs_chunk=dA_cumsum.stride(1),
634 stride_dA_cs_csize=dA_cumsum.stride(2),
644 - stride_seq_idx_seqlen=seq_idx.stride(0),
635 )
636 return states
637
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+21 -35
@@ -14,8 +14,7 @@ from vllm.triton_utils import triton
14
15 from .ssd_bmm import _bmm_chunk_fwd
16 from .ssd_chunk_scan import _chunk_scan_fwd
17 -from .ssd_chunk_state import (_chunk_cumsum_fwd, _chunk_state_fwd,
18 - chunk_state_varlen)
17 +from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd
18 from .ssd_state_passing import _state_passing_fwd
19
20 TRITON_22 = version.parse(triton.__version__) >= version.parse('2.2.0')
@@ -37,9 +36,9 @@ def _mamba_chunk_scan_combined_fwd(x,
36 dt_bias=None,
37 initial_states=None,
38 seq_idx=None,
40 - chunk_indices=None,
41 - chunk_offsets=None,
39 cu_seqlens=None,
40 + cu_chunk_seqlens=None,
41 + last_chunk_indices=None,
42 dt_softplus=False,
43 dt_limit=(0.0, float("inf")),
44 state_dtype=None):
@@ -56,7 +55,7 @@ def _mamba_chunk_scan_combined_fwd(x,
55 if D is not None:
56 assert D.shape == (nheads, headdim) or D.shape == (nheads, )
57 if seq_idx is not None:
59 - assert seq_idx.shape == (seqlen, )
58 + assert seq_idx.shape == (cu_chunk_seqlens.shape[0] - 1, )
59 if B.stride(-1) != 1:
60 B = B.contiguous()
61 if C.stride(-1) != 1:
@@ -89,6 +88,7 @@ def _mamba_chunk_scan_combined_fwd(x,
88 dA_cumsum, dt = _chunk_cumsum_fwd(dt,
89 A,
90 chunk_size,
91 + cu_chunk_seqlens,
92 dt_bias=dt_bias,
93 dt_softplus=dt_softplus,
94 dt_limit=dt_limit)
@@ -99,36 +99,31 @@ def _mamba_chunk_scan_combined_fwd(x,
99 x,
100 dt,
101 dA_cumsum,
102 - seq_idx=seq_idx,
102 + cu_chunk_seqlens,
103 states_in_fp32=True)
104
105 # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
106 # (middle term of factorization of off-diag blocks; A terms)
107 - # - for handling chunked prefill, this requires i) initial_states
108 - # ii) seq_idx iii) is_cont_batched and (iv) chunk_offsets to be all specified.
107 + # - for handling chunked prefill, this requires i) initial_states and
108 + # ii) seq_idx to be all specified.
109 # - When a new seq_idx is detected, we will stop passing the prev_state
110 # and switch accordingly to the init_state corresponding to the new seq_idx.
111 - # - We will also make sure that the dA_cumsum is taken only from the start of the
112 - # sequence (hence we need the full dA_cumsum tensor and not just the values at chunk boundaries)
113 - # - this will ensure that states will be updated with the rightmost flushed seq_idx
114 - # of the previous chunk. This implies that the first chunk of states is either 0
115 - # or equal to init_states of the first example.
111 states = _state_passing_fwd(
112 rearrange(states, "... p n -> ... (p n)"),
113 dA_cumsum, # (nheads, nchunks, chunk_size)
114 + cu_chunk_seqlens,
115 initial_states=rearrange(initial_states, "... p n -> ... (p n)")
116 if initial_states is not None else
117 None, # (batch, nheads, headdim*dstate)
118 seq_idx=seq_idx,
123 - out_dtype=state_dtype if state_dtype is not None else C.dtype,
124 - chunk_offsets=chunk_offsets)
119 + out_dtype=state_dtype if state_dtype is not None else C.dtype)
120 states = rearrange(states, "... (p n) -> ... p n", n=dstate)
121
122 # 4. Compute batched matrix multiply for C_j^T B_i terms
123 CB = _bmm_chunk_fwd(C,
124 B,
125 chunk_size,
131 - seq_idx=seq_idx,
126 + cu_chunk_seqlens,
127 output_dtype=torch.float32)
128
129 # 5. Scan and compute the diagonal blocks, taking into
@@ -148,26 +143,15 @@ def _mamba_chunk_scan_combined_fwd(x,
143 dA_cumsum,
144 C,
145 states,
146 + cu_chunk_seqlens,
147 out, # in-place update
148 seq_idx,
149 D=D,
150 z=z,
155 - chunk_indices=chunk_indices,
156 - chunk_offsets=chunk_offsets,
151 initial_states=initial_states,
152 )
153
160 - varlen_states = chunk_state_varlen(
161 - B,
162 - x,
163 - dt,
164 - dA_cumsum,
165 - cu_seqlens,
166 - states,
167 - initial_states=initial_states,
168 - )
169 -
170 - return varlen_states
154 + return states[last_chunk_indices]
155
156
157 def mamba_chunk_scan_combined_varlen(
@@ -178,14 +162,14 @@ def mamba_chunk_scan_combined_varlen(
162 C,
163 chunk_size,
164 cu_seqlens,
165 + cu_chunk_seqlens,
166 + last_chunk_indices,
167 seq_idx,
168 out,
169 D=None,
170 z=None,
171 dt_bias=None,
172 initial_states=None,
187 - chunk_indices=None,
188 - chunk_offsets=None,
173 dt_softplus=False,
174 dt_limit=(0.0, float("inf")),
175 state_dtype=None,
@@ -198,8 +182,10 @@ def mamba_chunk_scan_combined_varlen(
182 B: (seqlen, ngroups, dstate)
183 C: (seqlen, ngroups, dstate)
184 chunk_size: int
201 - seq_idx: (seqlen)
202 - cu_seqlens: (batch + 1)
185 + cu_seqlens: (batch + 1,)
186 + cu_chunk_seqlens: (nchunks + 1,)
187 + last_chunk_indices: (batch,)
188 + seq_idx: (nchunks,)
189 out: (seqlen, nheads, headdim) preallocated output tensor
190 D: (nheads, headdim) or (nheads,)
191 z: (seqlen, nheads, headdim)
@@ -228,9 +214,9 @@ def mamba_chunk_scan_combined_varlen(
214 dt_bias=dt_bias,
215 initial_states=initial_states,
216 seq_idx=seq_idx,
231 - chunk_indices=chunk_indices,
232 - chunk_offsets=chunk_offsets,
217 cu_seqlens=cu_seqlens,
218 + cu_chunk_seqlens=cu_chunk_seqlens,
219 + last_chunk_indices=last_chunk_indices,
220 dt_softplus=dt_softplus,
221 dt_limit=dt_limit,
222 state_dtype=state_dtype)
vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
+26 -71
@@ -30,8 +30,7 @@ def _state_passing_fwd_kernel(
30 dA_cs_ptr,
31 initstates_ptr,
32 seq_idx_ptr,
33 - chunk_offsets_ptr,
34 - chunk_meta_num,
33 + cu_chunk_seqlens_ptr,
34 # Matrix dimensions
35 dim: tl.constexpr,
36 nchunks,
@@ -50,94 +49,52 @@ def _state_passing_fwd_kernel(
49 stride_initstates_batch: tl.int64,
50 stride_initstates_head: tl.int64,
51 stride_initstates_dim: tl.constexpr,
53 - stride_seq_idx_seqlen: tl.constexpr,
52 + stride_seq_idx_chunk: tl.constexpr,
53 # Meta-parameters
54 HAS_INITSTATES: tl.constexpr,
55 BLOCK_SIZE: tl.constexpr,
56 ):
57 pid_h = tl.program_id(axis=1)
58 pid_m = tl.program_id(axis=0)
59 +
60 states_ptr += pid_h * stride_states_head
61 dA_cs_ptr += pid_h * stride_dA_cs_head + (chunk_size -
62 1) * stride_dA_cs_csize
63 out_ptr += pid_h * stride_out_head
64 - if HAS_INITSTATES:
65 - initstates_ptr += pid_h * stride_initstates_head
64
65 offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
66 states_ptrs = states_ptr + offs_m * stride_states_dim
67 out_ptrs = out_ptr + offs_m * stride_out_dim
68
71 - # - states will be the past state of the sequence that continues on the current check
72 - if not HAS_INITSTATES:
73 - states = tl.zeros((BLOCK_SIZE, ), dtype=tl.float32)
74 - else:
75 - initstates_ptr += offs_m * stride_initstates_dim
76 - initstates_ptrs = initstates_ptr
77 - # - for cont batches, for the first chunk mean it will be the first batch's
78 - # init state
69 + if HAS_INITSTATES:
70 + initstates_ptrs = initstates_ptr \
71 + + pid_h * stride_initstates_head \
72 + + offs_m * stride_initstates_dim
73 +
74 states = tl.load(initstates_ptrs, mask=offs_m < dim,
75 other=0.0).to(tl.float32)
76 + else:
77 + states = tl.zeros((BLOCK_SIZE, ), dtype=tl.float32)
78
82 - tl.store(out_ptrs, states, mask=offs_m < dim)
83 - out_ptrs += stride_out_chunk
84 - prev_seq_idx_chunk_end = 0
85 - logical_chunk_idx = 0
86 - for c in range(nchunks - 1):
79 + prev_seq_idx = 0
80 + for c in range(nchunks):
81 new_states = tl.load(states_ptrs, mask=offs_m < dim,
82 other=0.0).to(tl.float32)
83 dA_cs = tl.load(dA_cs_ptr).to(tl.float32)
90 - scale_mask = True
91 - # - the seq to pass forward is the one that is flushed to the right
92 - # boundary.
93 - # - that is given by seq_idx_chunk_end below: the sequence index at the end of the chunk.
94 - seq_idx_chunk_end = tl.load(seq_idx_ptr +
95 - (min((c + 1) * chunk_size, seqlen) - 1) *
96 - stride_seq_idx_seqlen)
97 -
98 - if HAS_INITSTATES:
99 - if prev_seq_idx_chunk_end != seq_idx_chunk_end:
100 - # this means in the current chunk the rightmost flushed seq
101 - # has changed.
102 - # - so we do not propagate the state from previous chunk
103 - # - but rather we load that sequence's init state
104 - initstates_ptrs = initstates_ptr + seq_idx_chunk_end * stride_initstates_batch
105 -
106 - # - update state with seq_idx_new's init state
84 + seq_idx = tl.load(seq_idx_ptr + c * stride_seq_idx_chunk)
85 + # we have started a new sequence
86 + if prev_seq_idx != seq_idx:
87 + if HAS_INITSTATES:
88 + initstates_ptrs = initstates_ptr + seq_idx * stride_initstates_batch \
89 + + pid_h * stride_initstates_head \
90 + + offs_m * stride_initstates_dim
91 states = tl.load(initstates_ptrs, mask=offs_m < dim,
92 other=0.0).to(tl.float32)
93 + else:
94 + states = tl.zeros((BLOCK_SIZE, ), dtype=tl.float32)
95
110 - # - we need to consider the cumsum only of the last sequence in the chunk
111 - # - find its starting position (given by c_off of the logical chunk index)
112 - # - and subtract the cumsum just before that position from the total cumsum
113 - # - first, update the logical chunk index (add the number of sequences in the current physical chunk):
114 - # sequence index at the start of the current chunk
115 - seq_idx_chunk_start = tl.load(seq_idx_ptr +
116 - min(c * chunk_size, seqlen) *
117 - stride_seq_idx_seqlen)
118 - logical_chunk_idx += seq_idx_chunk_end - seq_idx_chunk_start
119 - # - load the chunk offset:
120 - c_off = tl.load(chunk_offsets_ptr + logical_chunk_idx,
121 - mask=logical_chunk_idx < chunk_meta_num,
122 - other=0)
123 - # - if offset is 0, then the sequence starts at the beginning of the chunk, and we don't need to subtract anything
124 - if c_off > 0:
125 - # - dA_cs_ptr currently points to the cumsum at the end of the chunk - subtract the chunk size and add the offset
126 - dA_cs_boundary = tl.load(
127 - dA_cs_ptr - (chunk_size - 1) * stride_dA_cs_csize +
128 - (c_off - 1) * stride_dA_cs_csize,
129 - mask=(c_off - 1) > -1 and c_off < chunk_size,
130 - other=0.0)
131 - dA_cs -= dA_cs_boundary
132 -
133 - # - increment logical chunk index for every physical chunk
134 - logical_chunk_idx += 1
135 - else:
136 - scale_mask = seq_idx_chunk_end == prev_seq_idx_chunk_end
137 - prev_seq_idx_chunk_end = seq_idx_chunk_end
138 -
139 - scale = tl.where(scale_mask, tl.exp(dA_cs), 0.0)
140 - states = scale * states + new_states
96 + prev_seq_idx = seq_idx
97 + states = tl.exp(dA_cs) * states + new_states
98 tl.store(out_ptrs, states, mask=offs_m < dim)
99
100 states_ptrs += stride_states_chunk
@@ -148,8 +105,8 @@ def _state_passing_fwd_kernel(
105 def _state_passing_fwd(
106 states,
107 dA_cumsum,
108 + cu_chunk_seqlens,
109 seq_idx,
152 - chunk_offsets,
110 initial_states=None,
111 out_dtype=None,
112 ):
@@ -175,9 +132,7 @@ def _state_passing_fwd(
132 dA_cs_ptr=dA_cumsum,
133 initstates_ptr=initial_states,
134 seq_idx_ptr=seq_idx,
178 - chunk_offsets_ptr=chunk_offsets,
179 - chunk_meta_num=len(chunk_offsets)
180 - if chunk_offsets is not None else 0,
135 + cu_chunk_seqlens_ptr=cu_chunk_seqlens,
136 dim=dim,
137 nchunks=nchunks,
138 seqlen=seqlen if seq_idx is not None else 0,
@@ -194,7 +149,7 @@ def _state_passing_fwd(
149 stride_initstates_batch=initial_states_strides[0],
150 stride_initstates_head=initial_states_strides[1],
151 stride_initstates_dim=initial_states_strides[2],
197 - stride_seq_idx_seqlen=seq_idx.stride(0),
152 + stride_seq_idx_chunk=seq_idx.stride(0),
153 HAS_INITSTATES=initial_states is not None,
154 )
155 return out
vllm/model_executor/models/plamo2.py
+4 -4
@@ -260,9 +260,9 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
260 prep_initial_states = attn_metadata.prep_initial_states
261 chunk_size = attn_metadata.chunk_size
262 seq_idx_p = attn_metadata.seq_idx_p
263 - chunk_indices_p = attn_metadata.chunk_indices_p
264 - chunk_offsets_p = attn_metadata.chunk_offsets_p
263 query_start_loc_p = attn_metadata.query_start_loc_p
264 + cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
265 + last_chunk_indices_p = attn_metadata.last_chunk_indices_p
266
267 # 1. Gated MLP's linear projection
268 projected_states = self.in_proj(hidden_states)
@@ -368,9 +368,9 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
368 self.num_heads // self.tp_size, self.head_dim),
369 dt_bias=self.dt_bias,
370 seq_idx=seq_idx_p,
371 - chunk_indices=chunk_indices_p,
372 - chunk_offsets=chunk_offsets_p,
371 cu_seqlens=query_start_loc_p,
372 + cu_chunk_seqlens=cu_chunk_seqlen_p,
373 + last_chunk_indices=last_chunk_indices_p,
374 initial_states=initial_states,
375 dt_softplus=True,
376 dt_limit=(0.0, float("inf")),
vllm/v1/attention/backends/mamba2_attn.py
+80 -109
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import math
3 from dataclasses import dataclass
4 from typing import Optional
5
@@ -8,6 +7,7 @@ import torch
7
8 from vllm.attention.backends.abstract import AttentionBackend
9 from vllm.config import VllmConfig
10 +from vllm.utils import cdiv
11 from vllm.v1.attention.backends.mamba_attn import (
12 BaseMambaAttentionMetadataBuilder)
13 from vllm.v1.attention.backends.utils import (PAD_SLOT_ID,
@@ -17,91 +17,6 @@ from vllm.v1.attention.backends.utils import (PAD_SLOT_ID,
17 from vllm.v1.kv_cache_interface import AttentionSpec
18
19
20 -def _query_start_loc_to_chunk_indices_offsets(
21 - query_start_loc: torch.Tensor, chunk_size: int,
22 - total_seqlens: int) -> tuple[torch.Tensor, torch.Tensor]:
23 - """
24 - Args:
25 - query_start_loc (torch.Tensor): 1D tensor of cumulative sequence
26 - lengths, shape (num_seqs + 1,).
27 - The first element should be 0. Each entry represents the starting
28 - index of a sequence in the flattened token array.
29 - chunk_size (int): The size of each physical mamba chunk
30 - (number of tokens per chunk).
31 - total_seqlens (int): The total number of tokens in the batch.
32 -
33 - Returns:
34 - Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
35 - - chunk_indices (torch.Tensor): 1D tensor of indices
36 - indicating the physical chunk for each logical chunk.
37 - - chunk_offsets (torch.Tensor): 1D tensor of offsets
38 - indicating the starting index of each logical chunk within
39 - its physical chunk.
40 -
41 - This function computes the chunk indices and offsets for the given
42 - query_start_loc and chunk_size. Both are tensors of integers with length N,
43 - where N is the number of logical (pseudo) chunks.
44 - A logical chunk is a sequence of tokens that are all part of the same
45 - sequence and are all in the same physical mamba chunk.
46 - In other words, a logical chunk changes every time we cross a sequence
47 - boundary or a physical mamba chunk boundary.
48 - Logical chunks are needed to handle batched requests with initial states
49 - (see _state_passing_fwd and _chunk_scan_fwd).
50 - The chunk_indices tensor contains the index of the physical chunk for each
51 - logical chunk.
52 - The chunk_offsets tensor contains the offset (AKA starting index) of the
53 - logical chunk in the physical chunk.
54 -
55 - Example:
56 - query_start_loc = [0, 5, 10]
57 - chunk_size = 8
58 - total_seqlens = 10
59 - -> chunk_indices = [0, 0, 1]
60 - -> chunk_offsets = [0, 5, 0]
61 -
62 - In this example, we have 2 sequences, each with 5 tokens. The physical
63 - chunk size is 8 tokens.
64 - We have three logical chunks:
65 - - the first logical chunk starts at token 0 in the first physical chunk
66 - and contains all 5 tokens from the first sequence
67 - - the second logical chunk starts at token 5 in the first physical chunk
68 - and contains first 3 tokens from the second sequence
69 - - the third logical chunk starts at token 0 in the second physical chunk
70 - and contains the remaining 2 tokens from the second sequence
71 - """
72 -
73 - cu_seqlens = query_start_loc[1:] # remove prepended 0
74 -
75 - # outputs will have length expansion of chunks that do not divide
76 - # chunk_size
77 - N = math.ceil(total_seqlens / chunk_size) + (cu_seqlens[:-1] % chunk_size
78 - > 0).sum()
79 - chunk_indices = torch.arange(N,
80 - dtype=torch.int,
81 - device=query_start_loc.device)
82 - chunk_offsets = torch.zeros((N, ),
83 - dtype=torch.int,
84 - device=query_start_loc.device)
85 -
86 - p = 0 # num of insertions
87 - for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]):
88 -
89 - # if does not divide chunk_size, then there is one chunk insertion
90 - p += (s % chunk_size > 0)
91 -
92 - # get the dimensions
93 - # - the + 1 for _e is to shift the boundary by one chunk
94 - # - this shifting is not needed if chunk_size divides e
95 - _s, _e = s // chunk_size + p, e // chunk_size + p + (e % chunk_size
96 - > 0)
97 -
98 - # adjust indices and offsets
99 - chunk_indices[_s:_e] -= p
100 - chunk_offsets[_s] = s % chunk_size
101 -
102 - return chunk_indices, chunk_offsets
103 -
104 -
20 class Mamba2AttentionBackend(AttentionBackend):
21
22 @staticmethod
@@ -125,8 +40,16 @@ class Mamba2AttentionMetadata:
40 # the batch has no prefill request.
41 has_initial_states_p: Optional[torch.Tensor]
42 seq_idx_p: Optional[torch.Tensor]
128 - chunk_indices_p: Optional[torch.Tensor]
129 - chunk_offsets_p: Optional[torch.Tensor]
43 +
44 + # cu_chunk_seqlen_p is a tensor of shape (nchunks+1,) that contains, for
45 + # each chunk, its offests into the varlen sequence dimension. It is defined
46 + # such that the i-th chunk contains tokens from cu_chunk_seqlen_p[i] to
47 + # cu_chunk_seqlen_p[i+1].
48 + cu_chunk_seqlen_p: Optional[torch.Tensor]
49 +
50 + # last_chunk_indices_p is a tensor of shape (batch,) that contains the
51 + # index of the last chunk for every sequence in the (prefill) batch.
52 + last_chunk_indices_p: Optional[torch.Tensor]
53
54 state_indices_tensor: torch.Tensor # shape: [batch,]
55
@@ -151,13 +74,14 @@ class Mamba2AttentionMetadataBuilder(
74 common_attn_metadata: CommonAttentionMetadata,
75 fast_build: bool = False) -> Mamba2AttentionMetadata:
76 num_reqs = common_attn_metadata.num_reqs
154 - query_start_loc_p = None
77 seq_lens = common_attn_metadata.seq_lens
78
79 + query_start_loc_p = None
80 seq_idx_p = None
158 - chunk_indices_p, chunk_offsets_p = None, None
81 + cu_chunk_seqlen_p = None
82 + last_chunk_indices_p = None
83 +
84 # Need flags to indicate if there are initial states
160 - # currently we really only support the FlashAttention backend
85 has_initial_states_p = None
86 prep_initial_states = False
87
@@ -171,7 +95,7 @@ class Mamba2AttentionMetadataBuilder(
95 common_attn_metadata,
96 decode_threshold=self.reorder_batch_threshold))
97
174 - # Compute seq_idx, chunk_indices and chunk_offsets for prefill only
98 + # Compute seq_idx for prefill only
99 if num_prefills > 0:
100 #[batch,]
101 has_initial_states_cpu = (
@@ -184,21 +108,68 @@ class Mamba2AttentionMetadataBuilder(
108 query_start_loc_p = common_attn_metadata.query_start_loc[
109 -num_prefills - 1:] - num_decode_tokens
110
187 - seq_idx_p = torch.repeat_interleave(torch.arange(
188 - num_prefills,
189 - dtype=torch.int32,
190 - device=query_start_loc_p.device),
191 - query_start_loc_p.diff(),
192 - output_size=num_prefill_tokens)
193 -
194 - # We compute metadata for chunked prefill once at the top level
195 - # model forward and reuse them in mamba layers. If not needed,
196 - # they will be ignored inside mamba kernels.
197 - if prep_initial_states:
198 - chunk_indices_p, chunk_offsets_p = (
199 - _query_start_loc_to_chunk_indices_offsets(
200 - query_start_loc_p, self.chunk_size,
201 - num_prefill_tokens))
111 + num_computed_tokens_p = \
112 + common_attn_metadata.num_computed_tokens_cpu[
113 + num_reqs - num_prefills:num_reqs]
114 + query_start_loc_p_cpu = common_attn_metadata.query_start_loc_cpu[
115 + -num_prefills - 1:] - num_decode_tokens
116 +
117 + # The code below carefully constructs the chunks such that:
118 + # 1. Chunks contain tokens from a *single* sequence only.
119 + # 2. For every sequence, we are guaranteed that we can
120 + # retrieve the mamba state *every* chunk_size tokens.
121 + # Constraint (1) dramatically simplifies the mamba2 kernels.
122 + # Constraint (2) dramatically simplifies the implementation
123 + # of prefix caching for mamba2 (wip). We need to take care
124 + # of the interaction with chunked prefill in order to
125 + # satisfy constraint (2).
126 + # TODO (tdoublep): This code could probably be optimized.
127 + cu_chunk_seqlen = []
128 + seq_idx = []
129 + last_chunk_indices = []
130 + seqlen_pos = 0
131 + for req_idx in range(num_prefills):
132 + this_num_computed = num_computed_tokens_p[req_idx].item()
133 + this_new_tokens = query_start_loc_p_cpu[req_idx + 1].item(
134 + ) - query_start_loc_p_cpu[req_idx].item()
135 +
136 + # if computed tokens are not chunk-aligned, use the first
137 + # chunk to finish it off
138 + if this_num_computed % self.chunk_size != 0:
139 + seq_idx.append(req_idx)
140 + cu_chunk_seqlen.append(seqlen_pos)
141 + # how many tokens to finish the chunk?
142 + chunk_len = cdiv(this_num_computed, self.chunk_size
143 + ) * self.chunk_size - this_num_computed
144 + # we can only use at most this_new_tokens
145 + chunk_len = min(chunk_len, this_new_tokens)
146 + seqlen_pos += chunk_len
147 + this_new_tokens -= chunk_len
148 +
149 + n_chunks = cdiv(this_new_tokens, self.chunk_size)
150 + for chunk in range(n_chunks):
151 + seq_idx.append(req_idx)
152 + cu_chunk_seqlen.append(seqlen_pos)
153 + chunk_len = min(self.chunk_size, this_new_tokens)
154 + seqlen_pos += chunk_len
155 + this_new_tokens -= chunk_len
156 +
157 + assert this_new_tokens == 0
158 + last_chunk_indices.append(len(cu_chunk_seqlen) - 1)
159 +
160 + cu_chunk_seqlen.append(seqlen_pos)
161 +
162 + seq_idx_p = torch.as_tensor(seq_idx,
163 + device=query_start_loc_p.device,
164 + dtype=torch.int32)
165 + cu_chunk_seqlen_p = torch.as_tensor(
166 + cu_chunk_seqlen,
167 + device=query_start_loc_p.device,
168 + dtype=torch.int32)
169 + last_chunk_indices_p = torch.as_tensor(
170 + last_chunk_indices,
171 + device=query_start_loc_p.device,
172 + dtype=torch.int32)
173
174 nums_dict, batch_ptr, token_chunk_offset_ptr = \
175 compute_causal_conv1d_metadata(query_start_loc_p)
@@ -222,9 +193,9 @@ class Mamba2AttentionMetadataBuilder(
193 chunk_size=self.chunk_size,
194 has_initial_states_p=has_initial_states_p,
195 seq_idx_p=seq_idx_p,
225 - chunk_indices_p=chunk_indices_p,
226 - chunk_offsets_p=chunk_offsets_p,
196 state_indices_tensor=state_indices_tensor,
197 + cu_chunk_seqlen_p=cu_chunk_seqlen_p,
198 + last_chunk_indices_p=last_chunk_indices_p,
199 nums_dict=nums_dict,
200 batch_ptr=batch_ptr,
201 token_chunk_offset_ptr=token_chunk_offset_ptr,