@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 00b31a36a2
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 13, 22:00 UTC

[V1] [Hybrid] Mamba1 Automatic Prefix Caching (#26377)

Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>

Asaf Joseph Gardin committed Nov 2, 2025 at 14:16 UTC 00b31a36a2d0de6d197a473280b2304d482714af
16 files changed +442 -153
csrc/mamba/mamba_ssm/selective_scan.h
+7 -1
@@ -24,6 +24,8 @@ struct SSMParamsBase {
24 int64_t pad_slot_id;
25
26 bool delta_softplus;
27 + bool cache_enabled;
28 + int block_size;
29
30 index_t A_d_stride;
31 index_t A_dstate_stride;
@@ -46,8 +48,9 @@ struct SSMParamsBase {
48 index_t out_z_batch_stride;
49 index_t out_z_d_stride;
50 index_t ssm_states_batch_stride;
49 - index_t ssm_states_dim_stride;
51 + index_t ssm_states_dim_stride;
52 index_t ssm_states_dstate_stride;
53 + index_t cache_indices_stride;
54
55 // Common data pointers.
56 void *__restrict__ A_ptr;
@@ -66,6 +69,9 @@ struct SSMParamsBase {
69 void *__restrict__ cache_indices_ptr;
70 void *__restrict__ has_initial_state_ptr;
71
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 };
76
77
csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+113 -21
@@ -119,7 +119,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
119
120 const int* cache_indices = params.cache_indices_ptr == nullptr ? nullptr
121 : reinterpret_cast<int *>(params.cache_indices_ptr);
122 - const int cache_index = cache_indices == nullptr ? batch_id : cache_indices[batch_id];
122 + const int cache_index = cache_indices == nullptr ? batch_id : cache_indices[batch_id];
123 // cache_index == params.pad_slot_id is defined as padding, so we exit early
124 if (cache_index == params.pad_slot_id){
125 return;
@@ -133,9 +133,18 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
133 input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + sequence_start_index * params.B_batch_stride + group_id * params.B_group_stride;
134 weight_t *C = reinterpret_cast<weight_t *>(params.C_ptr) + dim_id * kNRows * params.C_d_stride;
135 input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + sequence_start_index * params.C_batch_stride + group_id * params.C_group_stride;
136 - typename Ktraits::state_t *ssm_states = reinterpret_cast<typename Ktraits::state_t *>(params.ssm_states_ptr) +
137 - cache_index * params.ssm_states_batch_stride +
138 - dim_id * kNRows * params.ssm_states_dim_stride;
136 +
137 + typename Ktraits::state_t *ssm_states;
138 + if (params.cache_enabled) {
139 + // APC mode: ssm_states points to the base, we'll use absolute cache slots later
140 + ssm_states = reinterpret_cast<typename Ktraits::state_t *>(params.ssm_states_ptr) +
141 + dim_id * kNRows * params.ssm_states_dim_stride;
142 + } else {
143 + // Non-APC mode: offset by cache_index as before
144 + ssm_states = reinterpret_cast<typename Ktraits::state_t *>(params.ssm_states_ptr) +
145 + cache_index * params.ssm_states_batch_stride +
146 + dim_id * kNRows * params.ssm_states_dim_stride;
147 + }
148
149 float D_val[kNRows] = {0};
150 if (params.D_ptr != nullptr) {
@@ -159,7 +168,22 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
168 // }
169
170 constexpr int kChunkSize = kNThreads * kNItems;
162 - const int n_chunks = (seqlen + 2048 - 1) / 2048;
171 +
172 + // 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;
175 +
176 + const int* batch_cache_indices = cache_indices != nullptr ?
177 + cache_indices + batch_id * params.cache_indices_stride : nullptr;
178 + const int* block_idx_first_scheduled = params.block_idx_first_scheduled_token_ptr != nullptr ?
179 + reinterpret_cast<const int*>(params.block_idx_first_scheduled_token_ptr) : nullptr;
180 + const int* block_idx_last_scheduled = params.block_idx_last_scheduled_token_ptr != nullptr ?
181 + reinterpret_cast<const int*>(params.block_idx_last_scheduled_token_ptr) : nullptr;
182 + const int* initial_state_idx = params.initial_state_idx_ptr != nullptr ?
183 + reinterpret_cast<const int*>(params.initial_state_idx_ptr) : nullptr;
184 +
185 + const size_t load_cache_slot = params.cache_enabled && batch_cache_indices != nullptr ? batch_cache_indices[initial_state_idx[batch_id]] : cache_index;
186 +
187 for (int chunk = 0; chunk < n_chunks; ++chunk) {
188 input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems];
189
@@ -219,7 +243,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
243 if constexpr (kIsVariableC) {
244 auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1;
245 load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
222 - smem_load_weight_C, (seqlen - chunk * kChunkSize) * (1 ));
246 + smem_load_weight_C, (seqlen - chunk * kChunkSize) * (1));
247 if constexpr (!kIsVariableB) {
248 #pragma unroll
249 for (int r = 0; r < kNRows; ++r) {
@@ -242,7 +266,6 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
266 for (int i = 0; i < kNItems; ++i) {
267 thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]),
268 !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]);
245 -
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);
@@ -250,8 +273,24 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
273 }
274 }
275 // Initialize running total
253 -
254 - scan_t running_prefix = chunk > 0 ? smem_running_prefix[state_idx + r * MAX_DSTATE] : make_float2(1.0, has_initial_state ? float(ssm_states[state_idx * params.ssm_states_dstate_stride]): 0.0);
276 + scan_t running_prefix;
277 + if (chunk > 0) {
278 + running_prefix = smem_running_prefix[state_idx + r * MAX_DSTATE];
279 + } else {
280 + // Load initial state
281 + if (params.cache_enabled && has_initial_state && batch_cache_indices != nullptr) {
282 + size_t state_offset = load_cache_slot * params.ssm_states_batch_stride +
283 + r * params.ssm_states_dim_stride +
284 + state_idx * params.ssm_states_dstate_stride;
285 + running_prefix = make_float2(1.0, float(ssm_states[state_offset]));
286 + } else if (has_initial_state) {
287 + // Non-APC mode: load from current batch position
288 + running_prefix = make_float2(1.0, float(ssm_states[state_idx * params.ssm_states_dstate_stride]));
289 + } else {
290 + // No initial state
291 + running_prefix = make_float2(1.0, 0.0);
292 + }
293 + }
294
295 SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
296 typename Ktraits::BlockScanT(smem_scan).InclusiveScan(
@@ -260,8 +299,25 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
299 // There's a syncthreads in the scan op, so we don't need to sync here.
300 // Unless there's only 1 warp, but then it's the same thread (0) reading and writing.
301 if (threadIdx.x == 0) {
263 - smem_running_prefix[state_idx] = prefix_op.running_prefix;
264 - if (chunk == n_chunks - 1) {
302 + smem_running_prefix[state_idx + r * MAX_DSTATE] = prefix_op.running_prefix;
303 +
304 + // Store state at the end of each chunk when cache is enabled
305 + if (params.cache_enabled && batch_cache_indices != nullptr) {
306 +
307 + size_t cache_slot;
308 + if (chunk == n_chunks - 1) {
309 + cache_slot = batch_cache_indices[block_idx_last_scheduled[batch_id]];
310 + } else {
311 + cache_slot = batch_cache_indices[block_idx_first_scheduled[batch_id] + chunk];
312 + }
313 +
314 + size_t state_offset = cache_slot * params.ssm_states_batch_stride +
315 + r * params.ssm_states_dim_stride +
316 + state_idx * params.ssm_states_dstate_stride;
317 +
318 + ssm_states[state_offset] = typename Ktraits::state_t(prefix_op.running_prefix.y);
319 + } else if (!params.cache_enabled && chunk == n_chunks - 1) {
320 + // Non-APC mode: store only final state at current batch position
321 ssm_states[state_idx * params.ssm_states_dstate_stride] = typename Ktraits::state_t(prefix_op.running_prefix.y);
322 }
323 }
@@ -274,7 +330,6 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
330 }
331 }
332 }
277 -
333 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;
335 __syncthreads();
@@ -346,7 +401,9 @@ template<typename input_t, typename weight_t, typename state_t>
401 void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
402
403 #ifndef USE_ROCM
349 - if (params.seqlen <= 128) {
404 + if (params.cache_enabled && params.block_size == 1024) {
405 + selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream);
406 + } else if (params.seqlen <= 128) {
407 selective_scan_fwd_launch<32, 4, input_t, weight_t, state_t>(params, stream);
408 } else if (params.seqlen <= 256) {
409 selective_scan_fwd_launch<32, 8, input_t, weight_t, state_t>(params, stream);
@@ -358,7 +415,9 @@ void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
415 selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream);
416 }
417 #else
361 - if (params.seqlen <= 256) {
418 + if (params.cache_enabled && params.block_size == 1024) {
419 + selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream);
420 + } else if (params.seqlen <= 256) {
421 selective_scan_fwd_launch<64, 4, input_t, weight_t, state_t>(params, stream);
422 } else if (params.seqlen <= 512) {
423 selective_scan_fwd_launch<64, 8, input_t, weight_t, state_t>(params, stream);
@@ -437,13 +496,17 @@ void set_ssm_params_fwd(SSMParamsBase &params,
496 const std::optional<at::Tensor>& D,
497 const std::optional<at::Tensor>& delta_bias,
498 const torch::Tensor ssm_states,
440 - bool has_z,
499 + bool has_z,
500 bool delta_softplus,
501 const std::optional<at::Tensor>& query_start_loc,
502 const std::optional<at::Tensor>& cache_indices,
503 const std::optional<at::Tensor>& has_initial_state,
504 bool varlen,
446 - int64_t pad_slot_id) {
505 + int64_t pad_slot_id,
506 + int64_t block_size,
507 + const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
508 + const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
509 + const std::optional<torch::Tensor> &initial_state_idx) {
510
511 // Reset the parameters
512 memset(&params, 0, sizeof(params));
@@ -477,6 +540,14 @@ void set_ssm_params_fwd(SSMParamsBase &params,
540 params.cache_indices_ptr = cache_indices.has_value() ? cache_indices.value().data_ptr() : nullptr;
541 params.has_initial_state_ptr = has_initial_state.has_value() ? has_initial_state.value().data_ptr() : nullptr;
542
543 + // Set cache parameters - cache is enabled if we have direct cache writing params
544 + params.cache_enabled = block_idx_first_scheduled_token.has_value();
545 + params.block_size = static_cast<int>(block_size);
546 +
547 + // Set direct cache writing pointers
548 + params.block_idx_first_scheduled_token_ptr = block_idx_first_scheduled_token.has_value() ? block_idx_first_scheduled_token.value().data_ptr() : nullptr;
549 + params.block_idx_last_scheduled_token_ptr = block_idx_last_scheduled_token.has_value() ? block_idx_last_scheduled_token.value().data_ptr() : nullptr;
550 + params.initial_state_idx_ptr = initial_state_idx.has_value() ? initial_state_idx.value().data_ptr() : nullptr;
551
552 // All stride are in elements, not bytes.
553 params.A_d_stride = A.stride(0);
@@ -504,9 +575,11 @@ void set_ssm_params_fwd(SSMParamsBase &params,
575 params.out_d_stride = out.stride(0);
576
577 params.ssm_states_batch_stride = ssm_states.stride(0);
507 - params.ssm_states_dim_stride = ssm_states.stride(1);
578 + params.ssm_states_dim_stride = ssm_states.stride(1);
579 params.ssm_states_dstate_stride = ssm_states.stride(2);
580
581 + params.cache_indices_stride = cache_indices.has_value() ? cache_indices.value().stride(0) : 0;
582 +
583 }
584 else{
585 if (!is_variable_B) {
@@ -537,8 +610,10 @@ void set_ssm_params_fwd(SSMParamsBase &params,
610 params.out_d_stride = out.stride(1);
611
612 params.ssm_states_batch_stride = ssm_states.stride(0);
540 - params.ssm_states_dim_stride = ssm_states.stride(1);
613 + params.ssm_states_dim_stride = ssm_states.stride(1);
614 params.ssm_states_dstate_stride = ssm_states.stride(2);
615 +
616 + params.cache_indices_stride = cache_indices.has_value() ? cache_indices.value().stride(0) : 0;
617 }
618 }
619
@@ -554,7 +629,11 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
629 const torch::Tensor &ssm_states,
630 // used to identify padding entries if cache_indices provided
631 // in case of padding, the kernel will return early
557 - int64_t pad_slot_id) {
632 + int64_t pad_slot_id,
633 + int64_t block_size,
634 + const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
635 + const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
636 + const std::optional<torch::Tensor> &initial_state_idx) {
637 auto input_type = u.scalar_type();
638 auto weight_type = A.scalar_type();
639 TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
@@ -646,7 +725,16 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
725 auto cache_indices_ = cache_indices.value();
726 TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int);
727 TORCH_CHECK(cache_indices_.is_cuda());
649 - CHECK_SHAPE(cache_indices_, batch_size);
728 +
729 + // cache_indices can be either 1D (batch_size,) for non-APC mode
730 + // or 2D (batch_size, max_positions) for APC mode
731 + const bool is_apc_mode = block_idx_first_scheduled_token.has_value();
732 + if (is_apc_mode) {
733 + TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
734 + TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
735 + } else {
736 + CHECK_SHAPE(cache_indices_, batch_size);
737 + }
738 }
739
740
@@ -686,7 +774,11 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
774 cache_indices,
775 has_initial_state,
776 varlen,
689 - pad_slot_id
777 + pad_slot_id,
778 + block_size,
779 + block_idx_first_scheduled_token,
780 + block_idx_last_scheduled_token,
781 + initial_state_idx
782 );
783
784
csrc/ops.h
+13 -11
@@ -321,17 +321,19 @@ void dynamic_per_token_scaled_fp8_quant(
321 torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scale,
322 std::optional<torch::Tensor> const& scale_ub);
323
324 -void selective_scan_fwd(const torch::Tensor& u, const torch::Tensor& delta,
325 - const torch::Tensor& A, const torch::Tensor& B,
326 - const torch::Tensor& C,
327 - const std::optional<torch::Tensor>& D_,
328 - const std::optional<torch::Tensor>& z_,
329 - const std::optional<torch::Tensor>& delta_bias_,
330 - bool delta_softplus,
331 - const std::optional<torch::Tensor>& query_start_loc,
332 - const std::optional<torch::Tensor>& cache_indices,
333 - const std::optional<torch::Tensor>& has_initial_state,
334 - const torch::Tensor& ssm_states, int64_t pad_slot_id);
324 +void selective_scan_fwd(
325 + const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A,
326 + const torch::Tensor& B, const torch::Tensor& C,
327 + const std::optional<torch::Tensor>& D_,
328 + const std::optional<torch::Tensor>& z_,
329 + const std::optional<torch::Tensor>& delta_bias_, bool delta_softplus,
330 + const std::optional<torch::Tensor>& query_start_loc,
331 + const std::optional<torch::Tensor>& cache_indices,
332 + const std::optional<torch::Tensor>& has_initial_state,
333 + const torch::Tensor& ssm_states, int64_t pad_slot_id, int64_t block_size,
334 + const std::optional<torch::Tensor>& block_idx_first_scheduled_token,
335 + const std::optional<torch::Tensor>& block_idx_last_scheduled_token,
336 + const std::optional<torch::Tensor>& initial_state_idx);
337
338 torch::Tensor dynamic_4bit_int_moe_cpu(
339 torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
csrc/torch_bindings.cpp
+5 -1
@@ -611,7 +611,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
611 "Tensor? cache_indices,"
612 "Tensor? has_initial_state,"
613 "Tensor! ssm_states,"
614 - "int pad_slot_id) -> ()");
614 + "int pad_slot_id,"
615 + "int block_size,"
616 + "Tensor? block_idx_first_scheduled_token,"
617 + "Tensor? block_idx_last_scheduled_token,"
618 + "Tensor? initial_state_idx) -> ()");
619 ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
620
621 // Hadamard transforms
tests/kernels/mamba/test_mamba_ssm.py
+15
@@ -179,6 +179,10 @@ def selective_scan_opcheck_fn(
179 has_initial_state=None,
180 ssm_states=None,
181 pad_slot_id=PAD_SLOT_ID,
182 + block_size=2048,
183 + block_idx_first_scheduled_token=None,
184 + block_idx_last_scheduled_token=None,
185 + initial_state_idx=None,
186 ):
187 """if return_last_state is True, returns (out, last_state)
188 last_state has shape (batch, dim, dstate).
@@ -223,6 +227,10 @@ def selective_scan_opcheck_fn(
227 has_initial_state,
228 ssm_states,
229 pad_slot_id,
230 + block_size,
231 + block_idx_first_scheduled_token,
232 + block_idx_last_scheduled_token,
233 + initial_state_idx,
234 ),
235 test_utils=["test_schema", "test_faketensor"],
236 )
@@ -338,6 +346,11 @@ def test_selective_scan(
346 has_initial_state=torch.ones(batch_size, device=u.device, dtype=torch.bool)
347 if c > 0
348 else None,
349 + pad_slot_id=PAD_SLOT_ID,
350 + block_size=2048,
351 + block_idx_first_scheduled_token=None,
352 + block_idx_last_scheduled_token=None,
353 + initial_state_idx=None,
354 )
355 outs.append(out)
356 if len(outs) > 1:
@@ -372,6 +385,7 @@ def test_selective_scan(
385 delta_bias=delta_bias,
386 delta_softplus=delta_softplus,
387 ssm_states=state,
388 + block_size=2048,
389 )
390
391
@@ -586,6 +600,7 @@ def test_selective_scan_varlen(
600 padded_state_indices,
601 has_initial_state,
602 prev_state,
603 + block_size=2048,
604 )
605
606
tests/models/language/generation/test_hybrid.py
+14 -20
@@ -19,6 +19,8 @@ pytestmark = pytest.mark.hybrid_model
19 # meaning that it will be used in all tests in this file
20 # The rest of the models will only be tested by test_models
21
22 +APC_MULTIPLY_BY = 300
23 +
24 SSM_MODELS = [
25 "state-spaces/mamba-130m-hf",
26 "tiiuae/falcon-mamba-tiny-dev",
@@ -380,7 +382,7 @@ def _get_vLLM_output(
382 return outs, vllm_model
383
384
383 -@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
385 +@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
386 @pytest.mark.parametrize("max_tokens", [64])
387 @pytest.mark.parametrize("n_repetitions", [2])
388 # If num_logprobs is set to -1, then the stringent version
@@ -410,10 +412,8 @@ def test_apc_single_prompt(
412 check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
413 )
414
413 - MULTIPLE = 300
414 -
415 # Sample prompts.
416 - generated_prompts = [MULTIPLE * example_prompts[0]]
416 + generated_prompts = [APC_MULTIPLY_BY * example_prompts[0]]
417
418 max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
419 vllm_runner_kwargs = _get_vllm_runner_params(
@@ -446,7 +446,7 @@ def test_apc_single_prompt(
446 )
447
448
449 -@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
449 +@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
450 @pytest.mark.parametrize("max_tokens", [64])
451 @pytest.mark.parametrize("n_repetitions", [2])
452 # If num_logprobs is set to -1, then the stringent version
@@ -476,10 +476,8 @@ def test_apc_single_prompt_block_align_alignment(
476 check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
477 )
478
479 - MULTIPLE = 300
480 -
479 # Sample prompts. This custom prompt is used, as it causes the most issues
482 - generated_prompts = ["The president of the United States is " * MULTIPLE]
480 + generated_prompts = ["The president of the United States is " * APC_MULTIPLY_BY]
481
482 max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
483 vllm_runner_kwargs = _get_vllm_runner_params(
@@ -528,7 +526,7 @@ def test_apc_single_prompt_block_align_alignment(
526 )
527
528
531 -@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
529 +@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
530 @pytest.mark.parametrize("max_tokens", [64])
531 @pytest.mark.parametrize("n_repetitions", [2])
532 # If num_logprobs is set to -1, then the stringent version
@@ -558,10 +556,8 @@ def test_apc_multiple_prompts_all_cached_outputs(
556 check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
557 )
558
561 - MULTIPLE = 300
562 -
559 # Sample prompts.
564 - generated_prompts = [MULTIPLE * prompt for prompt in example_prompts]
560 + generated_prompts = [APC_MULTIPLY_BY * prompt for prompt in example_prompts]
561
562 max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
563 vllm_runner_kwargs = _get_vllm_runner_params(
@@ -595,7 +591,7 @@ def test_apc_multiple_prompts_all_cached_outputs(
591 )
592
593
598 -@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
594 +@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
595 @pytest.mark.parametrize("max_tokens", [64])
596 @pytest.mark.parametrize("n_repetitions", [2])
597 # If num_logprobs is set to -1, then the stringent version
@@ -625,12 +621,12 @@ def test_apc_multiple_prompts_block_align_alignment(
621 check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
622 )
623
628 - MULTIPLE = 300
629 -
624 # Sample prompts. This custom prompt is used, as it causes the most issues
625 prompt_text = "The president of the United States is "
626 prompt_offsets = [0, 3, 7, 13, 17, 22, 25, 31]
633 - generated_prompts = [prompt_text[offset:] * MULTIPLE for offset in prompt_offsets]
627 + generated_prompts = [
628 + prompt_text[offset:] * APC_MULTIPLY_BY for offset in prompt_offsets
629 + ]
630
631 max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
632 vllm_runner_kwargs = _get_vllm_runner_params(
@@ -679,7 +675,7 @@ def test_apc_multiple_prompts_block_align_alignment(
675 )
676
677
682 -@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
678 +@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
679 @pytest.mark.parametrize("max_tokens", [64])
680 @pytest.mark.parametrize("n_repetitions", [2])
681 # If num_logprobs is set to -1, then the stringent version
@@ -709,10 +705,8 @@ def test_apc_multiple_prompts_partial_cached_outputs(
705 check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
706 )
707
712 - MULTIPLE = 300
713 -
708 # Sample prompts.
715 - generated_prompts = [MULTIPLE * prompt for prompt in example_prompts]
709 + generated_prompts = [APC_MULTIPLY_BY * prompt for prompt in example_prompts]
710
711 max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
712 vllm_runner_kwargs = _get_vllm_runner_params(
vllm/_custom_ops.py
+8
@@ -1719,6 +1719,10 @@ def selective_scan_fwd(
1719 has_initial_state: torch.Tensor | None,
1720 ssm_states: torch.Tensor,
1721 pad_slot_id: int,
1722 + block_size: int = 1024,
1723 + block_idx_first_scheduled_token: torch.Tensor | None = None,
1724 + block_idx_last_scheduled_token: torch.Tensor | None = None,
1725 + initial_state_idx: torch.Tensor | None = None,
1726 ):
1727 torch.ops._C.selective_scan_fwd(
1728 u,
@@ -1735,6 +1739,10 @@ def selective_scan_fwd(
1739 has_initial_state,
1740 ssm_states,
1741 pad_slot_id,
1742 + block_size,
1743 + block_idx_first_scheduled_token,
1744 + block_idx_last_scheduled_token,
1745 + initial_state_idx,
1746 )
1747
1748
vllm/config/model.py
+6
@@ -1483,6 +1483,12 @@ class ModelConfig:
1483 if chunk_size is None:
1484 # used by e.g. Mamba2, NemotronH, Zamba
1485 chunk_size = getattr(self.hf_text_config, "chunk_size", None)
1486 +
1487 + # Since Mamba1 does not have a chunk notion
1488 + # we use a default chunk size of 1024.
1489 + if chunk_size is None:
1490 + chunk_size = 2048
1491 +
1492 return chunk_size
1493
1494 def get_multimodal_config(self) -> MultiModalConfig:
vllm/model_executor/layers/mamba/mamba_mixer.py
+60 -31
@@ -241,18 +241,21 @@ class MambaMixer(MambaBase, CustomOp):
241 forward_context: ForwardContext = get_forward_context()
242 attn_metadata = forward_context.attn_metadata
243
244 + assert self.cache_config is not None
245 + mamba_block_size = self.cache_config.mamba_block_size
246 + prefix_caching_enabled = self.cache_config.enable_prefix_caching
247 +
248 if attn_metadata is not None:
249 assert isinstance(attn_metadata, dict)
250 attn_metadata = attn_metadata[self.prefix]
247 - mamba1_metadata = attn_metadata
248 - assert isinstance(mamba1_metadata, Mamba1AttentionMetadata)
249 - query_start_loc = mamba1_metadata.query_start_loc
250 - state_indices_tensor = mamba1_metadata.state_indices_tensor
251 + assert isinstance(attn_metadata, Mamba1AttentionMetadata)
252 + query_start_loc_p = attn_metadata.query_start_loc_p
253 + state_indices_tensor = attn_metadata.state_indices_tensor
254 self_kv_cache = self.kv_cache[forward_context.virtual_engine]
255 conv_state = self_kv_cache[0].transpose(-1, -2)
256 ssm_state = self_kv_cache[1]
254 - has_initial_states = mamba1_metadata.has_initial_states
255 - num_padded_decodes = mamba1_metadata.num_padded_decodes
257 + has_initial_states_p = attn_metadata.has_initial_states_p
258 + num_padded_decodes = attn_metadata.num_padded_decodes
259
260 # 1. Gated MLP's linear projection
261 projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)
@@ -279,12 +282,8 @@ class MambaMixer(MambaBase, CustomOp):
282 hidden_states_BC,
283 gate,
284 state_indices_tensor,
282 - query_start_loc,
283 - has_initial_states,
285 num_prefill_tokens,
285 - num_decode_tokens,
286 num_prefills,
287 - num_decodes,
287 num_padded_decodes,
288 )
289 hidden_states_BC_p = prefill_decode_split.hidden_states_BC_p
@@ -293,8 +292,34 @@ class MambaMixer(MambaBase, CustomOp):
292 gate_d = prefill_decode_split.gate_d
293 state_indices_tensor_p = prefill_decode_split.state_indices_tensor_p
294 state_indices_tensor_d = prefill_decode_split.state_indices_tensor_d
296 - query_start_loc_p = prefill_decode_split.query_start_loc_p
297 - has_initial_states_p = prefill_decode_split.has_initial_states_p
295 +
296 + if prefix_caching_enabled:
297 + block_idx_last_computed_token_d, block_idx_last_computed_token_p = (
298 + torch.split(
299 + attn_metadata.block_idx_last_computed_token,
300 + [num_decodes, num_prefills],
301 + dim=0,
302 + )
303 + )
304 + block_idx_last_scheduled_token_d, block_idx_last_scheduled_token_p = (
305 + torch.split(
306 + attn_metadata.block_idx_last_scheduled_token,
307 + [num_decodes, num_prefills],
308 + dim=0,
309 + )
310 + )
311 +
312 + block_idx_first_scheduled_token_p = (
313 + attn_metadata.block_idx_first_scheduled_token_p
314 + )
315 + num_computed_tokens_p = attn_metadata.num_computed_tokens_p
316 + else:
317 + block_idx_last_computed_token_d = None
318 + block_idx_last_computed_token_p = None
319 + block_idx_last_scheduled_token_d = None
320 + block_idx_last_scheduled_token_p = None
321 + block_idx_first_scheduled_token_p = None
322 + num_computed_tokens_p = None
323
324 ssm_outputs = []
325
@@ -309,6 +334,11 @@ class MambaMixer(MambaBase, CustomOp):
334 has_initial_state=has_initial_states_p,
335 cache_indices=state_indices_tensor_p,
336 query_start_loc=query_start_loc_p,
337 + block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,
338 + block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,
339 + initial_state_idx=block_idx_last_computed_token_p,
340 + num_computed_tokens=num_computed_tokens_p,
341 + block_size_to_align=mamba_block_size,
342 )
343 # 3. State Space Model sequence transformations.
344 discrete_time_step_p, B_p, C_p = self._ssm_transform(
@@ -331,10 +361,24 @@ class MambaMixer(MambaBase, CustomOp):
361 cache_indices=state_indices_tensor_p,
362 has_initial_state=has_initial_states_p,
363 query_start_loc=query_start_loc_p,
364 + block_size=mamba_block_size,
365 + block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,
366 + block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,
367 + initial_state_idx=block_idx_last_computed_token_p,
368 )
369 ssm_outputs.append(scan_out_p)
370
371 if has_decode:
372 + if prefix_caching_enabled:
373 + state_indices_tensor_d_input = state_indices_tensor_d.gather(
374 + 1, block_idx_last_computed_token_d.unsqueeze(1)
375 + ).squeeze(1)
376 + state_indices_tensor_d_output = state_indices_tensor_d.gather(
377 + 1, block_idx_last_scheduled_token_d.unsqueeze(1)
378 + ).squeeze(1)
379 + else:
380 + state_indices_tensor_d_input = state_indices_tensor_d
381 + state_indices_tensor_d_output = state_indices_tensor_d
382 # 2. Convolution sequence transformation
383 conv_out_d = causal_conv1d_update(
384 hidden_states_BC_d.transpose(0, 1),
@@ -343,6 +387,8 @@ class MambaMixer(MambaBase, CustomOp):
387 self.conv1d.bias,
388 self.activation,
389 conv_state_indices=state_indices_tensor_d,
390 + block_idx_last_scheduled_token=block_idx_last_scheduled_token_d,
391 + initial_state_idx=block_idx_last_computed_token_d,
392 ).transpose(0, 1)
393
394 # 3. State Space Model sequence transformation.
@@ -364,7 +410,8 @@ class MambaMixer(MambaBase, CustomOp):
410 gate_d.transpose(0, 1),
411 time_proj_bias,
412 dt_softplus=True,
367 - state_batch_indices=state_indices_tensor_d,
413 + state_batch_indices=state_indices_tensor_d_input,
414 + dst_state_batch_indices=state_indices_tensor_d_output,
415 out=scan_outputs_d,
416 )
417 scan_outputs_d = scan_outputs_d.transpose(0, 1)
@@ -423,20 +470,14 @@ class PrefillDecodeSplit(NamedTuple):
470 gate_d: torch.Tensor
471 state_indices_tensor_p: torch.Tensor
472 state_indices_tensor_d: torch.Tensor
426 - query_start_loc_p: torch.Tensor | None
427 - has_initial_states_p: torch.Tensor | None
473
474
475 def split_batch_to_prefill_and_decode(
476 hidden_states_BC: torch.Tensor,
477 gate: torch.Tensor,
478 state_indices_tensor: torch.Tensor,
434 - query_start_loc: torch.Tensor,
435 - has_initial_states: torch.Tensor | None,
479 num_prefill_tokens: int,
437 - num_decode_tokens: int,
480 num_prefills: int,
439 - num_decodes: int,
481 num_padded_decodes: int,
482 ) -> PrefillDecodeSplit:
483 num_actual_tokens = num_prefill_tokens + num_padded_decodes
@@ -457,16 +498,6 @@ def split_batch_to_prefill_and_decode(
498 [num_padded_decodes, num_prefills],
499 dim=0,
500 )
460 - query_start_loc_p = (
461 - query_start_loc[-num_prefills - 1 :] - num_padded_decodes
462 - if num_prefills > 0
463 - else None
464 - )
465 - has_initial_states_p = (
466 - has_initial_states[-num_prefills:]
467 - if (has_initial_states is not None and num_prefills > 0)
468 - else None
469 - )
501
502 return PrefillDecodeSplit(
503 hidden_states_BC_p=hidden_states_BC_p,
@@ -475,8 +506,6 @@ def split_batch_to_prefill_and_decode(
506 gate_d=gate_d,
507 state_indices_tensor_p=state_indices_tensor_p,
508 state_indices_tensor_d=state_indices_tensor_d,
478 - query_start_loc_p=query_start_loc_p,
479 - has_initial_states_p=has_initial_states_p,
509 )
510
511
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+23 -1
@@ -375,6 +375,10 @@ def selective_scan_fn(
375 cache_indices=None,
376 has_initial_state=None,
377 pad_slot_id=PAD_SLOT_ID,
378 + block_size=1024,
379 + block_idx_first_scheduled_token=None,
380 + block_idx_last_scheduled_token=None,
381 + initial_state_idx=None,
382 ) -> torch.Tensor:
383 """
384 u: (dim, total_length) for varlen or (batch, dim, seqlen)
@@ -397,7 +401,10 @@ def selective_scan_fn(
401 x.shape=(dim,17)
402 cache_indices: (batch) int32
403 A tensor with each cell is a correspondent
400 - input and output ssm_state index
404 + input and output ssm_state indices
405 + - Without APC: (batch,) - single state index per batch item
406 + - With APC: (batch, max_positions) - cache block indices for read/write
407 + Each non-zero value indicates a cache block to load from and/or write to.
408 has_initial_state: (batch) bool
409 A tensor populated with ones and zeros,
410 indicate if the ssm_state at the corresponding index should be
@@ -408,6 +415,17 @@ def selective_scan_fn(
415 that will not be processed,
416 for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
417 in this case, the kernel will not process entries at indices 0 and 3
418 + block_size: int
419 + The block size to align the cached states to
420 + block_idx_first_scheduled_token: (batch,), dtype int32
421 + The pointer into cache_indices, where the first
422 + cache block to be filled is located.
423 + block_idx_last_scheduled_token: (batch,), dtype int32
424 + The pointer into cache_indices, where the last cache block
425 + to be filled is located.
426 + initial_state_idx: (batch,), dtype int32
427 + The pointer into cache_indices, where the cache block
428 + containing the initial state is located.
429 returns
430 output: (dim, total_length) for varlen or (batch, dim, seqlen)
431 supports inplace replacement
@@ -448,6 +466,10 @@ def selective_scan_fn(
466 has_initial_state,
467 ssm_states,
468 pad_slot_id,
469 + block_size,
470 + block_idx_first_scheduled_token,
471 + block_idx_last_scheduled_token,
472 + initial_state_idx,
473 )
474
475 if z is None:
vllm/model_executor/models/config.py
+1 -1
@@ -299,7 +299,7 @@ class MambaModelConfig(VerifyAndUpdateConfig):
299 if model_config.supports_mamba_prefix_caching:
300 logger.info(
301 "Warning: Prefix caching is currently enabled. "
302 - "Its support for Mamba2 layers is experimental. "
302 + "Its support for Mamba layers is experimental. "
303 "Please report any issues you may observe."
304 )
305 else:
vllm/model_executor/models/jamba.py
+15 -6
@@ -38,7 +38,13 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38 from vllm.model_executor.models.llama import LlamaMLP as JambaMLP
39 from vllm.sequence import IntermediateTensors
40
41 -from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
41 +from .interfaces import (
42 + HasInnerState,
43 + IsHybrid,
44 + SupportsLoRA,
45 + SupportsMambaPrefixCaching,
46 + SupportsPP,
47 +)
48 from .utils import (
49 AutoWeightsLoader,
50 WeightsMapper,
@@ -454,7 +460,14 @@ class JambaModel(nn.Module):
460 return loaded_params
461
462
457 -class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid):
463 +class JambaForCausalLM(
464 + nn.Module,
465 + HasInnerState,
466 + SupportsLoRA,
467 + SupportsPP,
468 + IsHybrid,
469 + SupportsMambaPrefixCaching,
470 +):
471 hf_to_vllm_mapper = WeightsMapper(
472 orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"},
473 )
@@ -477,12 +490,8 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHyb
490
491 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
492 config = vllm_config.model_config.hf_config
480 - cache_config = vllm_config.cache_config
493 lora_config = vllm_config.lora_config
494 scheduler_config = vllm_config.scheduler_config
483 - assert not cache_config.enable_prefix_caching, (
484 - "Jamba currently does not support prefix caching"
485 - )
495
496 super().__init__()
497 self.config = config
vllm/model_executor/models/mamba.py
+4 -5
@@ -29,6 +29,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
29 from vllm.model_executor.models.interfaces import (
30 HasInnerState,
31 IsAttentionFree,
32 + SupportsMambaPrefixCaching,
33 SupportsPP,
34 )
35 from vllm.sequence import IntermediateTensors
@@ -193,15 +194,13 @@ class MambaModel(nn.Module):
194 return loaded_params
195
196
196 -class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
197 +class MambaForCausalLM(
198 + nn.Module, HasInnerState, IsAttentionFree, SupportsPP, SupportsMambaPrefixCaching
199 +):
200 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
201 config = vllm_config.model_config.hf_config
199 - cache_config = vllm_config.cache_config
202 lora_config = vllm_config.lora_config
203 self.scheduler_config = vllm_config.scheduler_config
202 - assert not cache_config.enable_prefix_caching, (
203 - "Mamba does not support prefix caching"
204 - )
204
205 super().__init__()
206 self.config = config
vllm/v1/attention/backends/mamba1_attn.py
+95 -16
@@ -7,11 +7,13 @@ import torch
7
8 from vllm.attention.backends.abstract import AttentionBackend
9 from vllm.attention.backends.utils import PAD_SLOT_ID
10 +from vllm.config import VllmConfig
11 from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder
12 from vllm.v1.attention.backends.utils import (
13 CommonAttentionMetadata,
14 split_decodes_and_prefills,
15 )
16 +from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
17
18
19 class Mamba1AttentionBackend(AttentionBackend):
@@ -22,32 +24,41 @@ class Mamba1AttentionBackend(AttentionBackend):
24
25 @dataclass
26 class Mamba1AttentionMetadata:
25 - query_start_loc: torch.Tensor
26 - context_lens_tensor: torch.Tensor
27 + query_start_loc_p: torch.Tensor
28 state_indices_tensor: torch.Tensor
28 - has_initial_states: torch.Tensor | None
29 + has_initial_states_p: torch.Tensor | None
30 num_prefills: int
31 num_prefill_tokens: int
32 num_decodes: int
33 num_decode_tokens: int
34 num_padded_decodes: int
35
36 + block_idx_last_scheduled_token: torch.Tensor # shape: [batch,]
37 + block_idx_first_scheduled_token_p: torch.Tensor # shape: [batch,]
38 + block_idx_last_computed_token: torch.Tensor # shape: [batch,]
39 + num_computed_tokens_p: torch.Tensor # shape: [batch,]
40 +
41
42 class Mamba1AttentionMetadataBuilder(
43 BaseMambaAttentionMetadataBuilder[Mamba1AttentionMetadata]
44 ):
45 + def __init__(
46 + self,
47 + kv_cache_spec: AttentionSpec,
48 + layer_names: list[str],
49 + vllm_config: VllmConfig,
50 + device: torch.device,
51 + ):
52 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
53 + assert isinstance(kv_cache_spec, MambaSpec)
54 +
55 def build(
56 self,
57 common_prefix_len: int,
58 common_attn_metadata: CommonAttentionMetadata,
59 fast_build: bool = False,
60 ) -> Mamba1AttentionMetadata:
45 - query_start_loc = common_attn_metadata.query_start_loc
46 -
47 - state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
48 - context_lens_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
49 - query_start_loc.device
50 - )
61 + num_reqs = common_attn_metadata.num_reqs
62
63 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
64 split_decodes_and_prefills(
@@ -55,32 +66,100 @@ class Mamba1AttentionMetadataBuilder(
66 )
67 )
68
58 - has_initial_states = None
69 + has_initial_states_p = None
70 + query_start_loc_p = None
71 padded_decodes = num_decodes
72 + num_computed_tokens, num_computed_tokens_p = None, None
73 + block_idx_first_scheduled_token = None
74 + block_idx_first_scheduled_token_p = None
75 +
76 + # TODO(@Josephasafg) Mamba1 and Mamba2 have a lot of code in common here.
77 + # We should consolidate this code
78 + if self.vllm_config.cache_config.enable_prefix_caching:
79 + # Return a tensor of shape (#requests, #max blocks)
80 + state_indices_tensor = common_attn_metadata.block_table_tensor
81 + mamba_block_size = self.kv_cache_spec.block_size
82 + num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
83 + self.device
84 + )
85 + (
86 + block_idx_last_computed_token,
87 + block_idx_first_scheduled_token,
88 + block_idx_last_scheduled_token,
89 + ) = self._compute_prefix_caching_block_indices(
90 + common_attn_metadata, mamba_block_size
91 + )
92 + else:
93 + # Always return just a single block per each request:
94 + state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
95 + block_idx_last_scheduled_token = None
96 + block_idx_last_computed_token = None
97
98 if num_prefills > 0:
62 - has_initial_states = context_lens_tensor > 0
99 + query_start_loc_p = (
100 + common_attn_metadata.query_start_loc[-num_prefills - 1 :]
101 + - num_decode_tokens
102 + )
103 + has_initial_states_cpu = (
104 + common_attn_metadata.num_computed_tokens_cpu[
105 + num_reqs - num_prefills : num_reqs
106 + ]
107 + > 0
108 + )
109 + has_initial_states_p = has_initial_states_cpu.to(
110 + common_attn_metadata.query_start_loc.device
111 + )
112 +
113 + if self.vllm_config.cache_config.enable_prefix_caching:
114 + assert num_computed_tokens is not None
115 + num_computed_tokens_p = num_computed_tokens[
116 + num_reqs - num_prefills : num_reqs
117 + ]
118 + assert block_idx_first_scheduled_token is not None
119 + block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[
120 + num_reqs - num_prefills : num_reqs
121 + ]
122 +
123 elif (
124 num_decodes > 0
125 and num_decodes <= self.decode_cudagraph_max_bs
126 and self.compilation_config.full_cuda_graph
127 ):
68 - state_indices_for_decode = state_indices_tensor[:num_decodes]
128 padded_decodes = self.vllm_config.pad_for_cudagraph(num_decodes)
129 self.state_indices_tensor[:num_decodes].copy_(
71 - state_indices_for_decode, non_blocking=True
130 + state_indices_tensor, non_blocking=True
131 )
132 state_indices_tensor = self.state_indices_tensor[:padded_decodes]
133 state_indices_tensor[num_decodes:] = PAD_SLOT_ID
134
135 + if self.vllm_config.cache_config.enable_prefix_caching:
136 + self.block_idx_last_scheduled_token[:num_decodes].copy_(
137 + block_idx_last_scheduled_token, non_blocking=True
138 + )
139 + block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
140 + :padded_decodes
141 + ]
142 + block_idx_last_scheduled_token[num_decodes:] = 0
143 +
144 + self.block_idx_last_computed_token[:num_decodes].copy_(
145 + block_idx_last_computed_token, non_blocking=True
146 + )
147 + block_idx_last_computed_token = self.block_idx_last_computed_token[
148 + :padded_decodes
149 + ]
150 + block_idx_last_computed_token[num_decodes:] = 0
151 +
152 return Mamba1AttentionMetadata(
77 - query_start_loc=query_start_loc,
78 - context_lens_tensor=context_lens_tensor,
79 - has_initial_states=has_initial_states,
153 + query_start_loc_p=query_start_loc_p,
154 + has_initial_states_p=has_initial_states_p,
155 state_indices_tensor=state_indices_tensor,
156 num_prefills=num_prefills,
157 num_prefill_tokens=num_prefill_tokens,
158 num_decodes=num_decodes,
159 num_decode_tokens=num_decode_tokens,
160 num_padded_decodes=padded_decodes,
161 + block_idx_last_scheduled_token=block_idx_last_scheduled_token,
162 + block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
163 + block_idx_last_computed_token=block_idx_last_computed_token,
164 + num_computed_tokens_p=num_computed_tokens_p,
165 )
vllm/v1/attention/backends/mamba2_attn.py
+6 -34
@@ -147,27 +147,6 @@ class Mamba2AttentionMetadataBuilder(
147 assert self.chunk_size is not None, (
148 "chunk_size needs to be set in the model config for Mamba2 models"
149 )
150 - if self.vllm_config.cache_config.enable_prefix_caching:
151 - self.state_indices_tensor = torch.empty(
152 - (
153 - self.decode_cudagraph_max_bs,
154 - cdiv(
155 - vllm_config.model_config.max_model_len, kv_cache_spec.block_size
156 - ),
157 - ),
158 - dtype=torch.int32,
159 - device=device,
160 - )
161 - self.block_idx_last_scheduled_token = torch.empty(
162 - (self.decode_cudagraph_max_bs,),
163 - dtype=torch.int32,
164 - device=device,
165 - )
166 - self.block_idx_last_computed_token = torch.empty(
167 - (self.decode_cudagraph_max_bs,),
168 - dtype=torch.int32,
169 - device=device,
170 - )
150
151 def build(
152 self,
@@ -202,20 +181,13 @@ class Mamba2AttentionMetadataBuilder(
181 num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
182 self.device
183 )
205 - # Block index of the last computed token
206 - block_idx_last_computed_token = (
207 - cdiv(num_computed_tokens, mamba_block_size) - 1
208 - )
209 - # which is <= block index for the first scheduled token
210 - block_idx_first_scheduled_token = (
211 - cdiv(num_computed_tokens + 1, mamba_block_size) - 1
212 - )
213 - # which is <= block index of the last scheduled token
214 - block_idx_last_scheduled_token = (
215 - cdiv(common_attn_metadata.seq_lens, mamba_block_size) - 1
184 + (
185 + block_idx_last_computed_token,
186 + block_idx_first_scheduled_token,
187 + block_idx_last_scheduled_token,
188 + ) = self._compute_prefix_caching_block_indices(
189 + common_attn_metadata, mamba_block_size
190 )
217 - # -1 in case it's non-computed and causes later issues with indexing
218 - block_idx_last_computed_token = block_idx_last_computed_token.clamp(min=0)
191 else:
192 # Always return just a single block per each request:
193 state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
vllm/v1/attention/backends/mamba_attn.py
+57 -5
@@ -7,6 +7,7 @@ from typing import ClassVar, TypeVar
7 import torch
8
9 from vllm.config import VllmConfig
10 +from vllm.utils.math_utils import cdiv
11 from vllm.v1.attention.backends.utils import (
12 AttentionCGSupport,
13 AttentionMetadataBuilder,
@@ -38,11 +39,35 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
39 self.vllm_config.scheduler_config.max_num_seqs,
40 self.compilation_config.max_cudagraph_capture_size,
41 )
41 - self.state_indices_tensor = torch.empty(
42 - (self.decode_cudagraph_max_bs,),
43 - dtype=torch.int32,
44 - device=device,
45 - )
42 +
43 + if self.vllm_config.cache_config.enable_prefix_caching:
44 + self.state_indices_tensor = torch.empty(
45 + (
46 + self.decode_cudagraph_max_bs,
47 + cdiv(
48 + self.vllm_config.model_config.max_model_len,
49 + self.kv_cache_spec.block_size,
50 + ),
51 + ),
52 + dtype=torch.int32,
53 + device=device,
54 + )
55 + self.block_idx_last_scheduled_token = torch.empty(
56 + (self.decode_cudagraph_max_bs,),
57 + dtype=torch.int32,
58 + device=device,
59 + )
60 + self.block_idx_last_computed_token = torch.empty(
61 + (self.decode_cudagraph_max_bs,),
62 + dtype=torch.int32,
63 + device=device,
64 + )
65 + else:
66 + self.state_indices_tensor = torch.empty(
67 + (self.decode_cudagraph_max_bs,),
68 + dtype=torch.int32,
69 + device=device,
70 + )
71
72 def build_for_cudagraph_capture(
73 self, common_attn_metadata: CommonAttentionMetadata
@@ -61,3 +86,30 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
86 m.max_query_len = 1 # decode-only
87
88 return self.build(0, m)
89 +
90 + def _compute_prefix_caching_block_indices(
91 + self,
92 + common_attn_metadata: CommonAttentionMetadata,
93 + mamba_block_size: int,
94 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
95 + num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
96 + self.device
97 + )
98 + # Block index of the last computed token
99 + block_idx_last_computed_token = cdiv(num_computed_tokens, mamba_block_size) - 1
100 + # which is <= block index for the first scheduled token
101 + block_idx_first_scheduled_token = (
102 + cdiv(num_computed_tokens + 1, mamba_block_size) - 1
103 + )
104 + # which is <= block index of the last scheduled token
105 + block_idx_last_scheduled_token = (
106 + cdiv(common_attn_metadata.seq_lens, mamba_block_size) - 1
107 + )
108 + # -1 in case it's non-computed and causes later issues with indexing
109 + block_idx_last_computed_token = block_idx_last_computed_token.clamp(min=0)
110 +
111 + return (
112 + block_idx_last_computed_token,
113 + block_idx_first_scheduled_token,
114 + block_idx_last_scheduled_token,
115 + )