@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bbf81f9a92

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

[Mamba1] - Kernel Level Chunk Alignment for Prefix Caching (#34798)

Signed-off-by: Josephasafg <ajgard7@gmail.com>

Asaf Gardin committed Mar 1, 2026 at 14:40 UTC bbf81f9a9284d572b69db2c4fb002c2a8a80d507

11 files changed +251 -146

csrc/mamba/mamba_ssm/selective_scan.h

+3 -1

@@ -17,7 +17,7 @@

17 struct SSMParamsBase {

18 using index_t = size_t;

20 - int batch, dim, seqlen, dstate, n_groups, n_chunks;

20 + int batch, dim, seqlen, dstate, n_groups;

21 int dim_ngroups_ratio;

22 bool is_variable_B;

23 bool is_variable_C;

@@ -72,6 +72,8 @@ struct SSMParamsBase {

72 void *__restrict__ block_idx_first_scheduled_token_ptr; // (batch,) - first block to write

73 void *__restrict__ block_idx_last_scheduled_token_ptr; // (batch,) - last block to write

74 void *__restrict__ initial_state_idx_ptr; // (batch,) - index of the initial state to use

75 + void *__restrict__ cu_chunk_seqlen_ptr; // (nchunks+1,) - cumulative chunk token offsets

76 + void *__restrict__ last_chunk_indices_ptr; // (batch,) - index of last chunk per sequence

77 };

csrc/mamba/mamba_ssm/selective_scan_fwd.cu

+68 -35

@@ -81,7 +81,6 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

81 constexpr bool kIsVariableC = Ktraits::kIsVariableC;

82 constexpr bool kHasZ = Ktraits::kHasZ;

83 constexpr bool kVarlen = Ktraits::kVarlen;

84 - constexpr int kNThreads = Ktraits::kNThreads;

84 constexpr int kNItems = Ktraits::kNItems;

85 constexpr int kNRows = Ktraits::kNRows;

86 constexpr bool kDirectIO = Ktraits::kDirectIO;

@@ -161,17 +160,8 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

160 }

161 }

162

164 -

165 - // for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {

166 - // smem_a[state_idx] = A[state_idx * params.A_dstate_stride];

167 - // smem_bc[state_idx] = B[state_idx * params.B_dstate_stride] * C[state_idx * params.C_dstate_stride];

168 - // }

169 -

170 - constexpr int kChunkSize = kNThreads * kNItems;

171 -

163 // Use block_size for chunking when APC is enabled, otherwise use 2048 for backwards compatibility

173 - const int iteration_chunk_size = params.cache_enabled ? params.block_size : 2048;

174 - const int n_chunks = (seqlen + iteration_chunk_size - 1) / iteration_chunk_size;

164 + const int block_size = params.cache_enabled ? params.block_size : 2048;

165

166 const int* batch_cache_indices = cache_indices != nullptr ?

167 cache_indices + batch_id * params.cache_indices_stride : nullptr;

@@ -181,10 +171,44 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

171 reinterpret_cast<const int*>(params.block_idx_last_scheduled_token_ptr) : nullptr;

172 const int* initial_state_idx = params.initial_state_idx_ptr != nullptr ?

173 reinterpret_cast<const int*>(params.initial_state_idx_ptr) : nullptr;

174 + const int* cu_chunk_seqlen = params.cu_chunk_seqlen_ptr != nullptr ?

175 + reinterpret_cast<const int*>(params.cu_chunk_seqlen_ptr) : nullptr;

176 + const int* last_chunk_indices = params.last_chunk_indices_ptr != nullptr ?

177 + reinterpret_cast<const int*>(params.last_chunk_indices_ptr) : nullptr;

178

179 const size_t load_cache_slot = params.cache_enabled && batch_cache_indices != nullptr ? batch_cache_indices[initial_state_idx[batch_id]] : cache_index;

180

181 + const int block_idx_first = (params.cache_enabled && block_idx_first_scheduled != nullptr) ?

182 + block_idx_first_scheduled[batch_id] : 0;

183 +

184 + // Determine chunk boundaries from pre-computed metadata (APC mode)

185 + // or fall back to simple block_size chunking.

186 + int first_chunk_idx, n_chunks;

187 + int current_position;

188 +

189 + if (cu_chunk_seqlen != nullptr && last_chunk_indices != nullptr) {

190 + const int last_chunk_idx = last_chunk_indices[batch_id];

191 + first_chunk_idx = (batch_id == 0) ? 0 : last_chunk_indices[batch_id - 1] + 1;

192 + n_chunks = last_chunk_idx - first_chunk_idx + 1;

193 + // Derive current_position: if the first chunk is partial (fills remainder

194 + // of a started block), offset into the block accordingly.

195 + const int first_chunk_tokens = cu_chunk_seqlen[first_chunk_idx + 1] - cu_chunk_seqlen[first_chunk_idx];

196 + const int chunk_start_offset = (n_chunks > 1 && first_chunk_tokens < block_size)

197 + ? (block_size - first_chunk_tokens) : 0;

198 + current_position = block_idx_first * block_size + chunk_start_offset;

199 + } else {

200 + first_chunk_idx = 0;

201 + n_chunks = (seqlen + block_size - 1) / block_size;

202 + current_position = 0;

203 + }

204 +

205 + int tokens_processed = 0;

206 +

207 for (int chunk = 0; chunk < n_chunks; ++chunk) {

208 + const int chunk_tokens = (cu_chunk_seqlen != nullptr)

209 + ? cu_chunk_seqlen[first_chunk_idx + chunk + 1] - cu_chunk_seqlen[first_chunk_idx + chunk]

210 + : min(block_size, seqlen - tokens_processed);

211 + if (chunk_tokens <= 0) break;

212 input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems];

213

214 __syncthreads();

@@ -193,12 +217,12 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

217 if constexpr (!kDirectIO) {

218 if (r > 0) { __syncthreads(); }

219 }

196 - load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, seqlen - chunk * kChunkSize);

220 + load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, chunk_tokens);

221 if constexpr (!kDirectIO) { __syncthreads(); }

198 - load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, seqlen - chunk * kChunkSize);

222 + load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, chunk_tokens);

223 }

200 - u += kChunkSize;

201 - delta += kChunkSize;

224 + u += chunk_tokens;

225 + delta += chunk_tokens;

226

227 float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems];

228 #pragma unroll

@@ -232,7 +256,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

256 weight_t B_vals[kNItems], C_vals[kNItems];

257 if constexpr (kIsVariableB) {

258 load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,

235 - smem_load_weight, (seqlen - chunk * kChunkSize) * (1));

259 + smem_load_weight, chunk_tokens);

260 if constexpr (!kIsVariableC) {

261 #pragma unroll

262 for (int r = 0; r < kNRows; ++r) {

@@ -243,7 +267,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

267 if constexpr (kIsVariableC) {

268 auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1;

269 load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,

246 - smem_load_weight_C, (seqlen - chunk * kChunkSize) * (1));

270 + smem_load_weight_C, chunk_tokens);

271 if constexpr (!kIsVariableB) {

272 #pragma unroll

273 for (int r = 0; r < kNRows; ++r) {

@@ -266,10 +290,8 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

290 for (int i = 0; i < kNItems; ++i) {

291 thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]),

292 !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]);

269 - if (seqlen % (kNItems * kNThreads) != 0) { // So that the last state is correct

270 - if (threadIdx.x * kNItems + i >= seqlen - chunk * kChunkSize) {

271 - thread_data[i] = make_float2(1.f, 0.f);

272 - }

293 + if (threadIdx.x * kNItems + i >= chunk_tokens) {

294 + thread_data[i] = make_float2(1.f, 0.f);

295 }

296 }

297 // Initialize running total

@@ -301,14 +323,14 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

323 if (threadIdx.x == 0) {

324 smem_running_prefix[state_idx + r * MAX_DSTATE] = prefix_op.running_prefix;

325

304 - // Store state at the end of each chunk when cache is enabled

326 + // Store state at the end of each aligned chunk when cache is enabled

327 if (params.cache_enabled && batch_cache_indices != nullptr) {

306 -

328 size_t cache_slot;

329 if (chunk == n_chunks - 1) {

330 cache_slot = batch_cache_indices[block_idx_last_scheduled[batch_id]];

331 } else {

311 - cache_slot = batch_cache_indices[block_idx_first_scheduled[batch_id] + chunk];

332 + const int block_idx_completed = (current_position + chunk_tokens - 1) / block_size;

333 + cache_slot = batch_cache_indices[block_idx_completed];

334 }

335

336 size_t state_offset = cache_slot * params.ssm_states_batch_stride +

@@ -331,38 +353,41 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {

353 }

354 }

355 input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + sequence_start_index * params.out_batch_stride

334 - + dim_id * kNRows * params.out_d_stride + chunk * kChunkSize;

356 + + dim_id * kNRows * params.out_d_stride + tokens_processed;

357 __syncthreads();

358 #pragma unroll

359 for (int r = 0; r < kNRows; ++r) {

360 if constexpr (!kDirectIO) {

361 if (r > 0) { __syncthreads(); }

362 }

341 - store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, seqlen - chunk * kChunkSize);

363 + store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, chunk_tokens);

364 }

365

366 if constexpr (kHasZ) {

367 input_t *z = reinterpret_cast<input_t *>(params.z_ptr) + sequence_start_index * params.z_batch_stride

346 - + dim_id * kNRows * params.z_d_stride + chunk * kChunkSize;

368 + + dim_id * kNRows * params.z_d_stride + tokens_processed;

369 input_t *out_z = reinterpret_cast<input_t *>(params.out_z_ptr) + sequence_start_index * params.out_z_batch_stride

348 - + dim_id * kNRows * params.out_z_d_stride + chunk * kChunkSize;

370 + + dim_id * kNRows * params.out_z_d_stride + tokens_processed;

371 #pragma unroll

372 for (int r = 0; r < kNRows; ++r) {

373 input_t z_vals[kNItems];

374 __syncthreads();

353 - load_input<Ktraits>(z + r * params.z_d_stride, z_vals, smem_load, seqlen - chunk * kChunkSize);

375 + load_input<Ktraits>(z + r * params.z_d_stride, z_vals, smem_load, chunk_tokens);

376 #pragma unroll

377 for (int i = 0; i < kNItems; ++i) {

378 float z_val = z_vals[i];

379 out_vals[r][i] *= z_val / (1 + expf(-z_val));

380 }

381 __syncthreads();

360 - store_output<Ktraits>(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, seqlen - chunk * kChunkSize);

382 + store_output<Ktraits>(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, chunk_tokens);

383 }

384 }

385

364 - Bvar += kChunkSize * 1;

365 - Cvar += kChunkSize * 1;

386 + Bvar += chunk_tokens;

387 + Cvar += chunk_tokens;

388 +

389 + tokens_processed += chunk_tokens;

390 + current_position += chunk_tokens;

391 }

392 }

393

@@ -506,7 +531,9 @@ void set_ssm_params_fwd(SSMParamsBase &params,

531 int64_t block_size,

532 const std::optional<torch::Tensor> &block_idx_first_scheduled_token,

533 const std::optional<torch::Tensor> &block_idx_last_scheduled_token,

509 - const std::optional<torch::Tensor> &initial_state_idx) {

534 + const std::optional<torch::Tensor> &initial_state_idx,

535 + const std::optional<torch::Tensor> &cu_chunk_seqlen,

536 + const std::optional<torch::Tensor> &last_chunk_indices) {

537

538 // Reset the parameters

539 memset(&params, 0, sizeof(params));

@@ -548,6 +575,8 @@ void set_ssm_params_fwd(SSMParamsBase &params,

575 params.block_idx_first_scheduled_token_ptr = block_idx_first_scheduled_token.has_value() ? block_idx_first_scheduled_token.value().data_ptr() : nullptr;

576 params.block_idx_last_scheduled_token_ptr = block_idx_last_scheduled_token.has_value() ? block_idx_last_scheduled_token.value().data_ptr() : nullptr;

577 params.initial_state_idx_ptr = initial_state_idx.has_value() ? initial_state_idx.value().data_ptr() : nullptr;

578 + params.cu_chunk_seqlen_ptr = cu_chunk_seqlen.has_value() ? cu_chunk_seqlen.value().data_ptr() : nullptr;

579 + params.last_chunk_indices_ptr = last_chunk_indices.has_value() ? last_chunk_indices.value().data_ptr() : nullptr;

580

581 // All stride are in elements, not bytes.

582 params.A_d_stride = A.stride(0);

@@ -633,7 +662,9 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,

662 int64_t block_size,

663 const std::optional<torch::Tensor> &block_idx_first_scheduled_token,

664 const std::optional<torch::Tensor> &block_idx_last_scheduled_token,

636 - const std::optional<torch::Tensor> &initial_state_idx) {

665 + const std::optional<torch::Tensor> &initial_state_idx,

666 + const std::optional<torch::Tensor> &cu_chunk_seqlen,

667 + const std::optional<torch::Tensor> &last_chunk_indices) {

668 auto input_type = u.scalar_type();

669 auto weight_type = A.scalar_type();

670 TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);

@@ -778,7 +809,9 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,

809 block_size,

810 block_idx_first_scheduled_token,

811 block_idx_last_scheduled_token,

781 - initial_state_idx

812 + initial_state_idx,

813 + cu_chunk_seqlen,

814 + last_chunk_indices

815 );

816

817

csrc/ops.h

+3 -1

@@ -371,7 +371,9 @@ void selective_scan_fwd(

371 const torch::Tensor& ssm_states, int64_t pad_slot_id, int64_t block_size,

372 const std::optional<torch::Tensor>& block_idx_first_scheduled_token,

373 const std::optional<torch::Tensor>& block_idx_last_scheduled_token,

374 - const std::optional<torch::Tensor>& initial_state_idx);

374 + const std::optional<torch::Tensor>& initial_state_idx,

375 + const std::optional<torch::Tensor>& cu_chunk_seqlen,

376 + const std::optional<torch::Tensor>& last_chunk_indices);

377

378 torch::Tensor dynamic_4bit_int_moe_cpu(

379 torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,

csrc/torch_bindings.cpp

+3 -1

@@ -640,7 +640,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

640 "int block_size,"

641 "Tensor? block_idx_first_scheduled_token,"

642 "Tensor? block_idx_last_scheduled_token,"

643 - "Tensor? initial_state_idx) -> ()");

643 + "Tensor? initial_state_idx,"

644 + "Tensor? cu_chunk_seqlen,"

645 + "Tensor? last_chunk_indices) -> ()");

646 ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);

647

648 // Hadamard transforms

tests/kernels/mamba/test_mamba_ssm.py

@@ -183,6 +183,8 @@ def selective_scan_opcheck_fn(

183 block_idx_first_scheduled_token=None,

184 block_idx_last_scheduled_token=None,

185 initial_state_idx=None,

186 + cu_chunk_seqlen=None,

187 + last_chunk_indices=None,

188 ):

189 """if return_last_state is True, returns (out, last_state)

190 last_state has shape (batch, dim, dstate).

@@ -231,6 +233,8 @@ def selective_scan_opcheck_fn(

233 block_idx_first_scheduled_token,

234 block_idx_last_scheduled_token,

235 initial_state_idx,

236 + cu_chunk_seqlen,

237 + last_chunk_indices,

238 ),

239 test_utils=["test_schema", "test_faketensor"],

240 )

vllm/_custom_ops.py

@@ -2021,6 +2021,8 @@ def selective_scan_fwd(

2021 block_idx_first_scheduled_token: torch.Tensor | None = None,

2022 block_idx_last_scheduled_token: torch.Tensor | None = None,

2023 initial_state_idx: torch.Tensor | None = None,

2024 + cu_chunk_seqlen: torch.Tensor | None = None,

2025 + last_chunk_indices: torch.Tensor | None = None,

2026 ):

2027 torch.ops._C.selective_scan_fwd(

2028 u,

@@ -2041,6 +2043,8 @@ def selective_scan_fwd(

2043 block_idx_first_scheduled_token,

2044 block_idx_last_scheduled_token,

2045 initial_state_idx,

2046 + cu_chunk_seqlen,

2047 + last_chunk_indices,

2048 )

2049

2050

vllm/model_executor/layers/mamba/mamba_mixer.py

@@ -271,6 +271,8 @@ class MambaMixer(MambaBase, PluggableLayer):

271 conv_state = self_kv_cache[0].transpose(-1, -2)

272 ssm_state = self_kv_cache[1]

273 has_initial_states_p = attn_metadata.has_initial_states_p

274 + cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p

275 + last_chunk_indices_p = attn_metadata.last_chunk_indices_p

276

277 # 1. Gated MLP's linear projection

278 projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)

@@ -376,6 +378,8 @@ class MambaMixer(MambaBase, PluggableLayer):

378 block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,

379 block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,

380 initial_state_idx=block_idx_last_computed_token_p,

381 + cu_chunk_seqlen=cu_chunk_seqlen_p,

382 + last_chunk_indices=last_chunk_indices_p,

383 )

384 ssm_outputs.append(scan_out_p)

385

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

@@ -497,6 +497,8 @@ def selective_scan_fn(

497 block_idx_first_scheduled_token=None,

498 block_idx_last_scheduled_token=None,

499 initial_state_idx=None,

500 + cu_chunk_seqlen=None,

501 + last_chunk_indices=None,

502 ) -> torch.Tensor:

503 """

504 u: (dim, total_length) for varlen or (batch, dim, seqlen)

@@ -588,6 +590,8 @@ def selective_scan_fn(

590 block_idx_first_scheduled_token,

591 block_idx_last_scheduled_token,

592 initial_state_idx,

593 + cu_chunk_seqlen,

594 + last_chunk_indices,

595 )

596

597 if z is None:

vllm/v1/attention/backends/mamba1_attn.py

+31 -2

@@ -1,9 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from dataclasses import dataclass

4 +from dataclasses import dataclass, replace

5 +from typing import Any

6 -from vllm.v1.attention.backend import AttentionBackend

7 +from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata

8 from vllm.v1.attention.backends.mamba_attn import (

9 BaseMambaAttentionMetadata,

10 BaseMambaAttentionMetadataBuilder,

@@ -29,3 +30,31 @@ class Mamba1AttentionMetadataBuilder(

30 BaseMambaAttentionMetadataBuilder[Mamba1AttentionMetadata]

31 ):

32 metadata_cls = Mamba1AttentionMetadata

33 +

34 + def build(

35 + self,

36 + common_prefix_len: int,

37 + common_attn_metadata: CommonAttentionMetadata,

38 + fast_build: bool = False,

39 + **kwargs: Any,

40 + ) -> Mamba1AttentionMetadata:

41 + common = self._compute_common_metadata(common_attn_metadata)

42 +

43 + if (

44 + common.num_prefills > 0

45 + and self.vllm_config.cache_config.mamba_cache_mode == "all"

46 + ):

47 + cu_chunk_seqlen_p, _, last_chunk_indices_p = (

48 + self._build_chunk_metadata_tensors(

49 + self.kv_cache_spec.block_size,

50 + common,

51 + common_attn_metadata,

52 + )

53 + )

54 + return replace(

55 + common,

56 + cu_chunk_seqlen_p=cu_chunk_seqlen_p,

57 + last_chunk_indices_p=last_chunk_indices_p,

58 + )

59 +

60 + return common

vllm/v1/attention/backends/mamba2_attn.py

+6 -106

@@ -7,7 +7,6 @@ from typing import Any

7 import torch

9 from vllm.config import VllmConfig

10 -from vllm.utils.math_utils import cdiv

10 from vllm.v1.attention.backend import (

11 AttentionBackend,

12 CommonAttentionMetadata,

@@ -105,14 +104,6 @@ class Mamba2AttentionMetadata(BaseMambaAttentionMetadata):

104

105 # Chunk-related metadata (only for prefill)

106 seq_idx_p: torch.Tensor | None = None

108 - # cu_chunk_seqlen_p is a tensor of shape (nchunks+1,) that contains, for

109 - # each chunk, its offsets into the varlen sequence dimension. It is defined

110 - # such that the i-th chunk contains tokens from cu_chunk_seqlen_p[i] to

111 - # cu_chunk_seqlen_p[i+1].

112 - cu_chunk_seqlen_p: torch.Tensor | None = None

113 - # last_chunk_indices_p is a tensor of shape (batch,) that contains the

114 - # index of the last chunk for every sequence in the (prefill) batch.

115 - last_chunk_indices_p: torch.Tensor | None = None

107

108

109 class Mamba2AttentionMetadataBuilder(

@@ -134,68 +125,6 @@ class Mamba2AttentionMetadataBuilder(

125 )

126 self.chunk_size: int = chunk_size

127

137 - def _compute_chunk_metadata(

138 - self,

139 - num_prefills: int,

140 - num_computed_tokens_p_cpu: torch.Tensor,

141 - query_start_loc_p_cpu: torch.Tensor,

142 - ) -> tuple[list[int], list[int], list[int]]:

143 - """

144 - Compute chunk-specific metadata for Mamba2.

145 -

146 - The code below carefully constructs the chunks such that:

147 - 1. Chunks contain tokens from a *single* sequence only.

148 - 2. For every sequence, we are guaranteed that we can

149 - retrieve the mamba state *every* chunk_size tokens.

150 - Constraint (1) dramatically simplifies the mamba2 kernels.

151 - Constraint (2) dramatically simplifies the implementation

152 - of prefix caching for mamba2 (wip). We need to take care

153 - of the interaction with chunked prefill in order to

154 - satisfy constraint (2).

155 - """

156 - # TODO (tdoublep): This code could probably be optimized.

157 - cu_chunk_seqlen = []

158 - seq_idx = []

159 - last_chunk_indices = []

160 - seqlen_pos = 0

161 -

162 - for req_idx in range(num_prefills):

163 - this_num_computed = num_computed_tokens_p_cpu[req_idx].item()

164 - this_new_tokens = (

165 - query_start_loc_p_cpu[req_idx + 1].item()

166 - - query_start_loc_p_cpu[req_idx].item()

167 - )

168 -

169 - # if computed tokens are not chunk-aligned, use the first

170 - # chunk to finish it off

171 - if this_num_computed % self.chunk_size != 0:

172 - seq_idx.append(req_idx)

173 - cu_chunk_seqlen.append(seqlen_pos)

174 - # how many tokens to finish the chunk?

175 - chunk_len = (

176 - cdiv(this_num_computed, self.chunk_size) * self.chunk_size

177 - - this_num_computed

178 - )

179 - # we can only use at most this_new_tokens

180 - chunk_len = min(chunk_len, this_new_tokens)

181 - seqlen_pos += chunk_len

182 - this_new_tokens -= chunk_len

183 -

184 - n_chunks = cdiv(this_new_tokens, self.chunk_size)

185 - for chunk in range(n_chunks):

186 - seq_idx.append(req_idx)

187 - cu_chunk_seqlen.append(seqlen_pos)

188 - chunk_len = min(self.chunk_size, this_new_tokens)

189 - seqlen_pos += chunk_len

190 - this_new_tokens -= chunk_len

191 -

192 - assert this_new_tokens == 0

193 - last_chunk_indices.append(len(cu_chunk_seqlen) - 1)

194 -

195 - cu_chunk_seqlen.append(seqlen_pos)

196 -

197 - return cu_chunk_seqlen, seq_idx, last_chunk_indices

198 -

128 def build(

129 self,

130 common_prefix_len: int,

@@ -220,41 +149,12 @@ class Mamba2AttentionMetadataBuilder(

149 else False

150 )

151

223 - num_reqs = common.num_reqs

224 - num_prefills = common.num_prefills

225 - num_decode_tokens = common.num_decode_tokens

226 -

227 - num_computed_tokens_cpu = (

228 - common_attn_metadata.compute_num_computed_tokens().cpu()

229 - )

230 - num_computed_tokens_p_cpu = num_computed_tokens_cpu[

231 - num_reqs - num_prefills : num_reqs

232 - ]

233 - query_start_loc_p_cpu = (

234 - common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]

235 - - num_decode_tokens

236 - )

237 -

238 - cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata(

239 - num_prefills,

240 - num_computed_tokens_p_cpu,

241 - query_start_loc_p_cpu,

242 - )

243 -

244 - seq_idx_p = torch.as_tensor(

245 - seq_idx,

246 - device=common_attn_metadata.query_start_loc.device,

247 - dtype=torch.int32,

248 - )

249 - cu_chunk_seqlen_p = torch.as_tensor(

250 - cu_chunk_seqlen,

251 - device=common_attn_metadata.query_start_loc.device,

252 - dtype=torch.int32,

253 - )

254 - last_chunk_indices_p = torch.as_tensor(

255 - last_chunk_indices,

256 - device=common_attn_metadata.query_start_loc.device,

257 - dtype=torch.int32,

152 + cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p = (

153 + self._build_chunk_metadata_tensors(

154 + self.chunk_size,

155 + common,

156 + common_attn_metadata,

157 + )

158 )

159

160 return replace(

vllm/v1/attention/backends/mamba_attn.py

+121

@@ -59,6 +59,15 @@ class BaseMambaAttentionMetadata:

59 # The following tensor is only used for prefix caching in align mode

60 seq_lens: torch.Tensor

62 + # cu_chunk_seqlen_p is a tensor of shape (nchunks+1,) that contains, for

63 + # each chunk, its offsets into the varlen sequence dimension. It is defined

64 + # such that the i-th chunk contains tokens from cu_chunk_seqlen_p[i] to

65 + # cu_chunk_seqlen_p[i+1].

66 + cu_chunk_seqlen_p: torch.Tensor | None = None

67 + # last_chunk_indices_p is a tensor of shape (batch,) that contains the

68 + # index of the last chunk for every sequence in the (prefill) batch.

69 + last_chunk_indices_p: torch.Tensor | None = None

70 +

71 # The following attributes are for triton implementation of causal_conv1d

72 nums_dict: dict | None = None

73 batch_ptr: torch.Tensor | None = None

@@ -185,6 +194,118 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

194 common_attn_metadata, num_accepted_tokens=num_accepted_tokens

195 )

196

197 + def _compute_chunk_metadata(

198 + self,

199 + chunk_size: int,

200 + num_prefills: int,

201 + num_computed_tokens_p_cpu: torch.Tensor,

202 + query_start_loc_p_cpu: torch.Tensor,

203 + ) -> tuple[list[int], list[int], list[int]]:

204 + """

205 + Compute chunk-specific metadata for Mamba models.

206 +

207 + The code below carefully constructs the chunks such that:

208 + 1. Chunks contain tokens from a *single* sequence only.

209 + 2. For every sequence, we are guaranteed that we can

210 + retrieve the mamba state *every* chunk_size tokens.

211 + Constraint (1) dramatically simplifies the mamba kernels.

212 + Constraint (2) dramatically simplifies the implementation

213 + of prefix caching for mamba (wip). We need to take care

214 + of the interaction with chunked prefill in order to

215 + satisfy constraint (2).

216 + """

217 + # TODO (tdoublep): This code could probably be optimized.

218 + cu_chunk_seqlen = []

219 + seq_idx = []

220 + last_chunk_indices = []

221 + seqlen_pos = 0

222 +

223 + for req_idx in range(num_prefills):

224 + this_num_computed = num_computed_tokens_p_cpu[req_idx].item()

225 + this_new_tokens = (

226 + query_start_loc_p_cpu[req_idx + 1].item()

227 + - query_start_loc_p_cpu[req_idx].item()

228 + )

229 +

230 + # if computed tokens are not chunk-aligned, use the first

231 + # chunk to finish it off

232 + if this_num_computed % chunk_size != 0:

233 + seq_idx.append(req_idx)

234 + cu_chunk_seqlen.append(seqlen_pos)

235 + # how many tokens to finish the chunk?

236 + chunk_len = (

237 + cdiv(this_num_computed, chunk_size) * chunk_size - this_num_computed

238 + )

239 + # we can only use at most this_new_tokens

240 + chunk_len = min(chunk_len, this_new_tokens)

241 + seqlen_pos += chunk_len

242 + this_new_tokens -= chunk_len

243 +

244 + n_chunks = cdiv(this_new_tokens, chunk_size)

245 + for chunk in range(n_chunks):

246 + seq_idx.append(req_idx)

247 + cu_chunk_seqlen.append(seqlen_pos)

248 + chunk_len = min(chunk_size, this_new_tokens)

249 + seqlen_pos += chunk_len

250 + this_new_tokens -= chunk_len

251 +

252 + assert this_new_tokens == 0

253 + last_chunk_indices.append(len(cu_chunk_seqlen) - 1)

254 +

255 + cu_chunk_seqlen.append(seqlen_pos)

256 +

257 + return cu_chunk_seqlen, seq_idx, last_chunk_indices

258 +

259 + def _build_chunk_metadata_tensors(

260 + self,

261 + chunk_size: int,

262 + common: M,

263 + common_attn_metadata: CommonAttentionMetadata,

264 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

265 + """

266 + Compute chunk metadata and return as device tensors.

267 + Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p).

268 + """

269 + num_reqs = common.num_reqs

270 + num_prefills = common.num_prefills

271 + num_decode_tokens = common.num_decode_tokens

272 +

273 + num_computed_tokens_cpu = (

274 + common_attn_metadata.compute_num_computed_tokens().cpu()

275 + )

276 + num_computed_tokens_p_cpu = num_computed_tokens_cpu[

277 + num_reqs - num_prefills : num_reqs

278 + ]

279 + query_start_loc_p_cpu = (

280 + common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]

281 + - num_decode_tokens

282 + )

283 +

284 + cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata(

285 + chunk_size,

286 + num_prefills,

287 + num_computed_tokens_p_cpu,

288 + query_start_loc_p_cpu,

289 + )

290 +

291 + device = common_attn_metadata.query_start_loc.device

292 + cu_chunk_seqlen_p = torch.as_tensor(

293 + cu_chunk_seqlen,

294 + device=device,

295 + dtype=torch.int32,

296 + )

297 + seq_idx_p = torch.as_tensor(

298 + seq_idx,

299 + device=device,

300 + dtype=torch.int32,

301 + )

302 + last_chunk_indices_p = torch.as_tensor(

303 + last_chunk_indices,

304 + device=device,

305 + dtype=torch.int32,

306 + )

307 + return cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p

308 +

309 def _compute_prefix_caching_block_indices(

310 self,

311 common_attn_metadata: CommonAttentionMetadata,