@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ea507c3a93
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

[V1] [Hybrid] Mamba2 Automatic Prefix Caching (#25752)

Signed-off-by: Stanislaw Wozniak <stw@zurich.ibm.com> Signed-off-by: Thomas Ortner <boh@zurich.ibm.com> Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Thomas Ortner <boh@zurich.ibm.com> Co-authored-by: Thomas Parnell <tpa@zurich.ibm.com>

Stan Wozniak committed Oct 4, 2025 at 06:34 UTC ea507c3a93feb66a7f4ff1905289c214f97701d6
18 files changed +917 -147
tests/models/language/generation/test_hybrid.py
+413 -1
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from typing import Callable
5 +
6 import pytest
7
8 from tests.models.registry import HF_EXAMPLE_MODELS
@@ -8,7 +10,7 @@ from tests.utils import multi_gpu_test
10 from vllm.engine.arg_utils import EngineArgs
11 from vllm.sampling_params import SamplingParams
12
11 -from ...utils import check_logprobs_close
13 +from ...utils import check_logprobs_close, check_outputs_equal
14
15 # Mark all tests as hybrid
16 pytestmark = pytest.mark.hybrid_model
@@ -332,3 +334,413 @@ def test_fp32_cache_state(
334 name_0="hf",
335 name_1="vllm",
336 )
337 +
338 +
339 +# Helper functions for the APC tests
340 +def _get_vllm_runner_params(model, max_model_len, tensor_parallel_size=1):
341 + return {
342 + 'model_name': model,
343 + 'enable_prefix_caching': False,
344 + 'max_model_len': max_model_len,
345 + 'tensor_parallel_size': tensor_parallel_size,
346 + 'gpu_memory_utilization': 0.4
347 + }
348 +
349 +
350 +def _get_vLLM_output(vllm_runner,
351 + kwargs,
352 + prompts,
353 + max_tokens,
354 + num_logprobs,
355 + num_repetitions=1,
356 + vllm_model=None):
357 + outs = []
358 + if vllm_model is None:
359 + vllm_model = vllm_runner(**kwargs)
360 + for _ in range(num_repetitions):
361 + if num_logprobs < 0:
362 + vllm_output = vllm_model.generate_greedy(prompts, max_tokens)
363 + else:
364 + vllm_output = vllm_model.generate_greedy_logprobs(
365 + prompts, max_tokens, num_logprobs)
366 + outs.append(vllm_output)
367 +
368 + return outs, vllm_model
369 +
370 +
371 +@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
372 +@pytest.mark.parametrize("max_tokens", [64])
373 +@pytest.mark.parametrize("n_repetitions", [2])
374 +# If num_logprobs is set to -1, then the stringent version
375 +# of the test is executed using `check_outputs_equal`
376 +# instead of `check_logprobs_close`
377 +@pytest.mark.parametrize("num_logprobs", [5])
378 +@pytest.mark.parametrize("tensor_parallel_size", [1])
379 +def test_apc_single_prompt(
380 + hf_runner,
381 + vllm_runner,
382 + example_prompts,
383 + monkeypatch,
384 + model: str,
385 + max_tokens: int,
386 + n_repetitions: int,
387 + num_logprobs: int,
388 + tensor_parallel_size: int,
389 +) -> None:
390 +
391 + try:
392 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
393 + model_info.check_available_online(on_fail="skip")
394 + model_info.check_transformers_version(on_fail="skip")
395 + except ValueError:
396 + pass
397 +
398 + compare_operator: Callable = check_logprobs_close \
399 + if num_logprobs > 0 else check_outputs_equal # type: ignore
400 +
401 + MULTIPLE = 300
402 +
403 + # Sample prompts.
404 + generated_prompts = [MULTIPLE * example_prompts[0]]
405 +
406 + max_model_len = max(
407 + len(prompt) + max_tokens for prompt in generated_prompts)
408 + vllm_runner_kwargs = _get_vllm_runner_params(
409 + model, max_model_len, tensor_parallel_size=tensor_parallel_size)
410 + vllm_runner_kwargs['mamba_ssm_cache_dtype'] = "float32"
411 + vllm_outputs_no_cache, _ = _get_vLLM_output(vllm_runner,
412 + vllm_runner_kwargs,
413 + generated_prompts, max_tokens,
414 + num_logprobs)
415 +
416 + vllm_runner_kwargs['enable_prefix_caching'] = True
417 + vllm_outputs_cache_rep, _ = _get_vLLM_output(vllm_runner,
418 + vllm_runner_kwargs,
419 + generated_prompts, max_tokens,
420 + num_logprobs, n_repetitions)
421 +
422 + for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
423 + # In the first repetition, the caches are filled
424 + # In the second repetition, these caches are reused
425 +
426 + compare_operator(
427 + outputs_0_lst=vllm_outputs_no_cache[0],
428 + outputs_1_lst=vllm_outputs_cache_itn,
429 + name_0="vllm_no_cache",
430 + name_1=f"vllm_cache_it_{r_idx + 1}",
431 + )
432 +
433 +
434 +@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
435 +@pytest.mark.parametrize("max_tokens", [64])
436 +@pytest.mark.parametrize("n_repetitions", [2])
437 +# If num_logprobs is set to -1, then the stringent version
438 +# of the test is executed using `check_outputs_equal`
439 +# instead of `check_logprobs_close`
440 +@pytest.mark.parametrize("num_logprobs", [5])
441 +@pytest.mark.parametrize("tensor_parallel_size", [1])
442 +def test_apc_single_prompt_block_align_alignment(
443 + hf_runner,
444 + vllm_runner,
445 + example_prompts,
446 + monkeypatch,
447 + model: str,
448 + max_tokens: int,
449 + n_repetitions: int,
450 + num_logprobs: int,
451 + tensor_parallel_size: int,
452 +) -> None:
453 +
454 + try:
455 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
456 + model_info.check_available_online(on_fail="skip")
457 + model_info.check_transformers_version(on_fail="skip")
458 + except ValueError:
459 + pass
460 +
461 + compare_operator: Callable = check_logprobs_close \
462 + if num_logprobs > 0 else check_outputs_equal # type: ignore
463 +
464 + MULTIPLE = 300
465 +
466 + # Sample prompts. This custom prompt is used, as it causes the most issues
467 + generated_prompts = ["The president of the United States is " * MULTIPLE]
468 +
469 + max_model_len = max(
470 + len(prompt) + max_tokens for prompt in generated_prompts)
471 + vllm_runner_kwargs = _get_vllm_runner_params(
472 + model, max_model_len, tensor_parallel_size=tensor_parallel_size)
473 + vllm_runner_kwargs['mamba_ssm_cache_dtype'] = "float32"
474 +
475 + vllm_outputs_no_cache, _ = _get_vLLM_output(vllm_runner,
476 + vllm_runner_kwargs,
477 + generated_prompts, max_tokens,
478 + num_logprobs)
479 +
480 + vllm_runner_kwargs['enable_prefix_caching'] = True
481 + with vllm_runner(**vllm_runner_kwargs) as vllm_model:
482 + # Retrieve the default mamba state block size
483 + mamba_block_size = vllm_model.llm.llm_engine.cache_config. \
484 + mamba_block_size
485 +
486 + # In case the hybrid model does not have the
487 + # "mamba_block_size" assume a fixed constant
488 + if mamba_block_size is None:
489 + mamba_block_size = 512
490 +
491 + mamba_block_size_multiplier = 10
492 + for offsets in [
493 + -3, 3, mamba_block_size // 4 + 3, mamba_block_size // 2 - 3
494 + ]:
495 +
496 + vllm_runner_kwargs[
497 + 'max_num_batched_tokens'] = mamba_block_size_multiplier * \
498 + mamba_block_size - offsets
499 + vllm_outputs_cache_rep, _ = _get_vLLM_output(vllm_runner,
500 + vllm_runner_kwargs,
501 + generated_prompts,
502 + max_tokens, num_logprobs,
503 + n_repetitions)
504 +
505 + # Check alignment of the output logits when using APC
506 + for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
507 + # In the first repetition, the caches are filled
508 + # In the second repetition, these caches are reused
509 +
510 + compare_operator(
511 + outputs_0_lst=vllm_outputs_no_cache[0],
512 + outputs_1_lst=vllm_outputs_cache_itn,
513 + name_0="vllm_no_cache",
514 + name_1=f"vllm_cache_it_{r_idx + 1}",
515 + )
516 +
517 +
518 +@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
519 +@pytest.mark.parametrize("max_tokens", [64])
520 +@pytest.mark.parametrize("n_repetitions", [2])
521 +# If num_logprobs is set to -1, then the stringent version
522 +# of the test is executed using `check_outputs_equal`
523 +# instead of `check_logprobs_close`
524 +@pytest.mark.parametrize("num_logprobs", [5])
525 +@pytest.mark.parametrize("tensor_parallel_size", [1])
526 +def test_apc_multiple_prompts_all_cached_outputs(
527 + hf_runner,
528 + vllm_runner,
529 + example_prompts,
530 + monkeypatch,
531 + model: str,
532 + max_tokens: int,
533 + n_repetitions: int,
534 + num_logprobs: int,
535 + tensor_parallel_size: int,
536 +) -> None:
537 +
538 + try:
539 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
540 + model_info.check_available_online(on_fail="skip")
541 + model_info.check_transformers_version(on_fail="skip")
542 + except ValueError:
543 + pass
544 +
545 + compare_operator: Callable = check_logprobs_close \
546 + if num_logprobs > 0 else check_outputs_equal # type: ignore
547 +
548 + MULTIPLE = 300
549 +
550 + # Sample prompts.
551 + generated_prompts = [MULTIPLE * prompt for prompt in example_prompts]
552 +
553 + max_model_len = max(
554 + len(prompt) + max_tokens for prompt in generated_prompts)
555 + vllm_runner_kwargs = _get_vllm_runner_params(
556 + model, max_model_len, tensor_parallel_size=tensor_parallel_size)
557 + vllm_runner_kwargs['mamba_ssm_cache_dtype'] = "float32"
558 +
559 + vllm_outputs_no_cache, _ = _get_vLLM_output(vllm_runner,
560 + vllm_runner_kwargs,
561 + generated_prompts, max_tokens,
562 + num_logprobs)
563 +
564 + vllm_runner_kwargs['enable_prefix_caching'] = True
565 + vllm_outputs_cache_rep, _ = _get_vLLM_output(vllm_runner,
566 + vllm_runner_kwargs,
567 + generated_prompts, max_tokens,
568 + num_logprobs, n_repetitions)
569 +
570 + for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
571 + # In the first repetition, the caches are filled
572 + # In the second repetition, these caches are reused
573 +
574 + compare_operator(
575 + outputs_0_lst=vllm_outputs_no_cache[0],
576 + outputs_1_lst=vllm_outputs_cache_itn,
577 + name_0="vllm_no_cache",
578 + name_1=f"vllm_cache_it_{r_idx + 1}",
579 + )
580 +
581 +
582 +@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
583 +@pytest.mark.parametrize("max_tokens", [64])
584 +@pytest.mark.parametrize("n_repetitions", [2])
585 +# If num_logprobs is set to -1, then the stringent version
586 +# of the test is executed using `check_outputs_equal`
587 +# instead of `check_logprobs_close`
588 +@pytest.mark.parametrize("num_logprobs", [5])
589 +@pytest.mark.parametrize("tensor_parallel_size", [1])
590 +def test_apc_multiple_prompts_block_align_alignment(
591 + hf_runner,
592 + vllm_runner,
593 + example_prompts,
594 + monkeypatch,
595 + model: str,
596 + max_tokens: int,
597 + n_repetitions: int,
598 + num_logprobs: int,
599 + tensor_parallel_size: int,
600 +) -> None:
601 +
602 + try:
603 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
604 + model_info.check_available_online(on_fail="skip")
605 + model_info.check_transformers_version(on_fail="skip")
606 + except ValueError:
607 + pass
608 +
609 + compare_operator: Callable = check_logprobs_close \
610 + if num_logprobs > 0 else check_outputs_equal # type: ignore
611 +
612 + MULTIPLE = 300
613 +
614 + # Sample prompts. This custom prompt is used, as it causes the most issues
615 + prompt_text = "The president of the United States is "
616 + prompt_offsets = [0, 3, 7, 13, 17, 22, 25, 31]
617 + generated_prompts = [
618 + prompt_text[offset:] * MULTIPLE for offset in prompt_offsets
619 + ]
620 +
621 + max_model_len = max(
622 + len(prompt) + max_tokens for prompt in generated_prompts)
623 + vllm_runner_kwargs = _get_vllm_runner_params(model, max_model_len,
624 + tensor_parallel_size)
625 + vllm_runner_kwargs['mamba_ssm_cache_dtype'] = "float32"
626 +
627 + vllm_outputs_no_cache, _ = _get_vLLM_output(vllm_runner,
628 + vllm_runner_kwargs,
629 + generated_prompts, max_tokens,
630 + num_logprobs)
631 +
632 + vllm_runner_kwargs['enable_prefix_caching'] = True
633 + with vllm_runner(**vllm_runner_kwargs) as vllm_model:
634 + # Retrieve the default mamba state block size
635 + mamba_block_size = vllm_model.llm.llm_engine.cache_config. \
636 + mamba_block_size
637 +
638 + # In case the hybrid model does not have the
639 + # "mamba_block_size" assume a fixed constant
640 + if mamba_block_size is None:
641 + mamba_block_size = 512
642 +
643 + mamba_block_size_multiplier = 10
644 + for offsets in [
645 + -3, 3, mamba_block_size // 4 + 3, mamba_block_size // 2 - 3
646 + ]:
647 +
648 + vllm_runner_kwargs[
649 + 'max_num_batched_tokens'] = mamba_block_size_multiplier * \
650 + mamba_block_size - offsets
651 + vllm_outputs_cache_rep, _ = _get_vLLM_output(vllm_runner,
652 + vllm_runner_kwargs,
653 + generated_prompts,
654 + max_tokens, num_logprobs,
655 + n_repetitions)
656 +
657 + # Check alignment of the output logits when using APC
658 + for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
659 + # In the first repetition, the caches are filled
660 + # In the second repetition, these caches are reused
661 +
662 + compare_operator(
663 + outputs_0_lst=vllm_outputs_no_cache[0],
664 + outputs_1_lst=vllm_outputs_cache_itn,
665 + name_0="vllm_no_cache",
666 + name_1=f"vllm_cache_it_{r_idx + 1}",
667 + )
668 +
669 +
670 +@pytest.mark.parametrize("model", [HYBRID_MODELS[3]])
671 +@pytest.mark.parametrize("max_tokens", [64])
672 +@pytest.mark.parametrize("n_repetitions", [2])
673 +# If num_logprobs is set to -1, then the stringent version
674 +# of the test is executed using `check_outputs_equal`
675 +# instead of `check_logprobs_close`
676 +@pytest.mark.parametrize("num_logprobs", [5])
677 +@pytest.mark.parametrize("tensor_parallel_size", [1])
678 +def test_apc_multiple_prompts_partial_cached_outputs(
679 + hf_runner,
680 + vllm_runner,
681 + example_prompts,
682 + monkeypatch,
683 + model: str,
684 + max_tokens: int,
685 + n_repetitions: int,
686 + num_logprobs: int,
687 + tensor_parallel_size: int,
688 +) -> None:
689 +
690 + try:
691 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
692 + model_info.check_available_online(on_fail="skip")
693 + model_info.check_transformers_version(on_fail="skip")
694 + except ValueError:
695 + pass
696 +
697 + compare_operator: Callable = check_logprobs_close \
698 + if num_logprobs > 0 else check_outputs_equal # type: ignore
699 +
700 + MULTIPLE = 300
701 +
702 + # Sample prompts.
703 + generated_prompts = [MULTIPLE * prompt for prompt in example_prompts]
704 +
705 + max_model_len = max(
706 + len(prompt) + max_tokens for prompt in generated_prompts)
707 + vllm_runner_kwargs = _get_vllm_runner_params(
708 + model, max_model_len, tensor_parallel_size=tensor_parallel_size)
709 + vllm_runner_kwargs['mamba_ssm_cache_dtype'] = "float32"
710 +
711 + vllm_outputs_no_cache, _ = _get_vLLM_output(vllm_runner,
712 + vllm_runner_kwargs,
713 + generated_prompts, max_tokens,
714 + num_logprobs)
715 +
716 + # Cache only part of all the prompts
717 + vllm_runner_kwargs['enable_prefix_caching'] = True
718 + vllm_outputs_partial_cache, vllm_model = _get_vLLM_output(
719 + vllm_runner, vllm_runner_kwargs, generated_prompts[:3], max_tokens,
720 + num_logprobs)
721 +
722 + compare_operator(
723 + outputs_0_lst=vllm_outputs_no_cache[0][:3],
724 + outputs_1_lst=vllm_outputs_partial_cache[0],
725 + name_0="vllm_no_cache",
726 + name_1="vllm_partial_cache",
727 + )
728 +
729 + vllm_outputs_cache_rep, _ = _get_vLLM_output(vllm_runner,
730 + vllm_runner_kwargs,
731 + generated_prompts,
732 + max_tokens,
733 + num_logprobs,
734 + n_repetitions,
735 + vllm_model=vllm_model)
736 +
737 + for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
738 + # In the first repetition, the caches are filled
739 + # In the second repetition, these caches are reused
740 +
741 + compare_operator(
742 + outputs_0_lst=vllm_outputs_no_cache[0],
743 + outputs_1_lst=vllm_outputs_cache_itn,
744 + name_0="vllm_no_cache",
745 + name_1=f"vllm_cache_it_{r_idx + 1}",
746 + )
vllm/config/cache.py
+2 -1
@@ -92,7 +92,8 @@ class CacheConfig:
92 mamba_page_size_padded: Optional[int] = None
93 """ Optional override for mamba page size; used by hybrid mamba/attention
94 models to ensure exact alignment with attention page size."""
95 -
95 + mamba_block_size: Optional[int] = None
96 + """Size of a contiguous cache block in number of tokens for mamba cache."""
97 mamba_cache_dtype: MambaDType = "auto"
98 """The data type to use for the Mamba cache (both the conv as well as the
99 ssm state). If set to 'auto', the data type will be inferred from the model
vllm/engine/arg_utils.py
+6 -1
@@ -1563,7 +1563,12 @@ class EngineArgs:
1563 self.enable_prefix_caching = False
1564
1565 if self.enable_prefix_caching is None:
1566 - self.enable_prefix_caching = True
1566 + # Disable prefix caching default for hybrid models
1567 + # since the feature is still experimental.
1568 + if model_config.is_hybrid:
1569 + self.enable_prefix_caching = False
1570 + else:
1571 + self.enable_prefix_caching = True
1572 else:
1573
1574 pooling_type = model_config.pooler_config.pooling_type
vllm/model_executor/layers/mamba/mamba_mixer2.py
+105 -8
@@ -489,6 +489,9 @@ class MambaMixer2(MambaBase, CustomOp):
489 # stay the same and reused for all mamba layers in the same iteration
490 attn_metadata: AttentionMetadata = forward_context.attn_metadata
491
492 + assert self.cache_config is not None
493 + mamba_block_size = self.cache_config.mamba_block_size
494 + prefix_caching_enabled = self.cache_config.enable_prefix_caching
495 if attn_metadata is not None:
496 assert isinstance(attn_metadata, dict)
497 attn_metadata = attn_metadata[self.prefix]
@@ -573,6 +576,25 @@ class MambaMixer2(MambaBase, CustomOp):
576 dim=0,
577 )
578
579 + if prefix_caching_enabled:
580 + # If prefix caching is enabled, retrieve the relevant variables
581 + # for prefill and decode
582 + last_state_idx_d, last_state_idx_p = torch.split(
583 + attn_metadata.last_state_idx, [num_decodes, num_prefills],
584 + dim=0)
585 + current_last_idx_d, current_last_idx_p = torch.split(
586 + attn_metadata.current_last_idx, [num_decodes, num_prefills],
587 + dim=0)
588 + # Prefill-only variables:
589 + current_first_idx_p = attn_metadata.current_first_idx_p
590 + context_lens_p = attn_metadata.context_lens_p
591 + last_computed_offset_p = attn_metadata.last_computed_offset_p
592 + else:
593 + last_state_idx_d, last_state_idx_p = None, None
594 + current_last_idx_d, current_last_idx_p = None, None
595 + current_first_idx_p = None
596 + context_lens_p = None
597 +
598 # Preallocate output tensor to avoid memcpy cost for merging prefill
599 # and decode outputs
600 preallocated_ssm_out = torch.empty(
@@ -592,8 +614,17 @@ class MambaMixer2(MambaBase, CustomOp):
614 # Process prefill requests
615 if has_prefill:
616 # 2. Convolution sequence transformation
595 - # - "cache_indices" updates the conv_state cache in positions
596 - # pointed to by "state_indices_tensor"
617 + # - It will read the initial states for every sequence,
618 + # that has "has_initial_states_p" == True,
619 + # from "cache_indices", using "state_indices_tensor_p".
620 + # - It updates the "conv_state" cache in positions pointed
621 + # to by "state_indices_tensor_p".
622 + # In particular, it will always write the state at the
623 + # sequence end.
624 + # In addition, "current_first_idx_p" and "current_last_idx_p"
625 + # are provided (which are pointers into
626 + # "state_indices_tensor_p"), it will write additional cache
627 + # states aligned at "block_size_to_align".
628 x = hidden_states_B_C_p.transpose(
629 0, 1) # this is the form that causal-conv see
630 hidden_states_B_C_p = causal_conv1d_fn(
@@ -604,6 +635,11 @@ class MambaMixer2(MambaBase, CustomOp):
635 conv_states=conv_state,
636 has_initial_state=has_initial_states_p,
637 cache_indices=state_indices_tensor_p,
638 + current_first_idx=current_first_idx_p,
639 + current_last_idx=current_last_idx_p,
640 + initial_state_idx=last_state_idx_p,
641 + context_lens=context_lens_p,
642 + block_size_to_align=mamba_block_size,
643 metadata=attn_metadata,
644 query_start_loc=query_start_loc_p).transpose(
645 0, 1)[:num_prefill_tokens]
@@ -614,9 +650,13 @@ class MambaMixer2(MambaBase, CustomOp):
650 # 3. State Space Model sequence transformation
651 initial_states = None
652 if (has_initial_states_p is not None and prep_initial_states):
653 + kernel_ssm_indices = state_indices_tensor_p
654 + if prefix_caching_enabled:
655 + kernel_ssm_indices = state_indices_tensor_p.gather(
656 + 1, last_state_idx_p.unsqueeze(1)).squeeze(1)
657 initial_states = torch.where(
658 has_initial_states_p[:, None, None, None],
619 - ssm_state[state_indices_tensor_p], 0)
659 + ssm_state[kernel_ssm_indices], 0)
660
661 # NOTE: final output is an in-place update of out tensor
662 varlen_states = mamba_chunk_scan_combined_varlen(
@@ -638,18 +678,71 @@ class MambaMixer2(MambaBase, CustomOp):
678 cu_chunk_seqlens=cu_chunk_seqlen_p,
679 last_chunk_indices=last_chunk_indices_p,
680 initial_states=initial_states,
681 + return_intermediate_states=prefix_caching_enabled,
682 dt_softplus=True,
683 dt_limit=(0.0, float("inf")),
684 out=preallocated_ssm_out_p.view(num_prefill_tokens, -1,
685 self.head_dim),
686 state_dtype=ssm_state.dtype)
687
647 - # update ssm states
648 - # - varlen state is a (num_prefills, nheads, headdim, dstate) tensor
649 - ssm_state[state_indices_tensor_p] = varlen_states
688 + if prefix_caching_enabled:
689 + # Save states for sequences with more than just the final state:
690 + n_blocks_to_fill = current_last_idx_p - current_first_idx_p
691 + for seq_idx in (n_blocks_to_fill > 0).nonzero().squeeze(1):
692 + cache_blocks_to_fill = state_indices_tensor_p[
693 + seq_idx, current_first_idx_p[seq_idx]:
694 + current_first_idx_p[seq_idx] +
695 + n_blocks_to_fill[seq_idx]]
696 + # chunks = [0 1 2 3 4 5 6 ...]
697 + # First aligned chunk would typically be:
698 + # mamba_block_size = 1024, chunk_size = 256
699 + # 1024 // 256 - 1 --> chunks[3]
700 + # But when last chunk wasn't block aligned:
701 + # - last_computed_offset_p[seq_idx] // chunk_size
702 + # e.g. 1000 // 256 -> 3 completed --> store chunk[0]
703 + # e.g. 513 // 256 -> 2 completed --> store chunk[1] (skip 1)
704 + # e.g. 256 // 256 -> 1 completed --> store chunk[2] (skip 2)
705 + # e.g. 10 // 256 -> 0 completed --> store chunk[3] (skip 3)
706 + chunk_stride = mamba_block_size // chunk_size
707 + first_aligned_chunk = \
708 + torch.concat([torch.zeros(1, \
709 + dtype=last_chunk_indices_p.dtype, \
710 + device=last_chunk_indices_p.device), \
711 + last_chunk_indices_p + 1])[seq_idx] \
712 + + chunk_stride - 1 \
713 + - last_computed_offset_p[seq_idx] // chunk_size
714 + from_where = varlen_states[
715 + first_aligned_chunk:first_aligned_chunk +
716 + n_blocks_to_fill[seq_idx] * chunk_stride:chunk_stride]
717 + ssm_state[cache_blocks_to_fill] = from_where
718 +
719 + #For all seqs, store the last state (Note: might be partial):
720 + ssm_state[state_indices_tensor_p.gather(1,
721 + current_last_idx_p.unsqueeze(1)).squeeze(1)] = \
722 + varlen_states[last_chunk_indices_p]
723 + else:
724 + # update ssm states
725 + # - varlen state is a (num_prefills, nheads, headdim, dstate)
726 + # tensor
727 + ssm_state[state_indices_tensor_p] = varlen_states
728
729 # Process decode requests
730 if has_decode:
731 + if prefix_caching_enabled:
732 + state_indices_tensor_d_input = \
733 + state_indices_tensor_d.gather(1,
734 + last_state_idx_d.unsqueeze(1)).squeeze(1)
735 + state_indices_tensor_d_output = \
736 + state_indices_tensor_d.gather(1,
737 + current_last_idx_d.unsqueeze(1)).squeeze(1)
738 + #Note:
739 + # for decode always: current_first_idx_d == current_last_idx_d
740 + # at block boundaries: current_first_idx_d > last_state_idx_d
741 + else:
742 + # Without caching, read and write in-place to the same blocks:
743 + state_indices_tensor_d_input = state_indices_tensor_d
744 + state_indices_tensor_d_output = state_indices_tensor_d
745 +
746 # 2. Convolution sequence transformation
747 hidden_states_B_C_d = causal_conv1d_update(
748 hidden_states_B_C_d,
@@ -657,7 +750,10 @@ class MambaMixer2(MambaBase, CustomOp):
750 conv_weights,
751 self.conv1d.bias,
752 self.activation,
660 - conv_state_indices=state_indices_tensor_d)
753 + conv_state_indices=state_indices_tensor_d,
754 + current_last_idx=current_last_idx_d,
755 + initial_state_idx=last_state_idx_d,
756 + )
757
758 hidden_states_d, B_d, C_d = split_hidden_states_B_C_fn(
759 hidden_states_B_C_d)
@@ -689,7 +785,8 @@ class MambaMixer2(MambaBase, CustomOp):
785 z=None,
786 dt_bias=dt_bias,
787 dt_softplus=True,
692 - state_batch_indices=state_indices_tensor_d,
788 + state_batch_indices=state_indices_tensor_d_input,
789 + dst_state_batch_indices=state_indices_tensor_d_output,
790 out=preallocated_ssm_out_d.view(num_decodes, -1,
791 self.head_dim),
792 )
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+174 -68
@@ -20,19 +20,23 @@ def _causal_conv1d_fwd_kernel( # continuous batching
20 w_ptr, # (dim, width)
21 bias_ptr,
22 initial_states_ptr, # conv_states_ptr
23 - cache_indices_ptr, # conv_state_indices_ptr
23 + cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains
24 + # the block indices relevant for each sequence
25 + # plus potential 0-padding at the beginning and at the end
26 has_initial_states_ptr,
27 query_start_loc_ptr,
28 batch_ptr,
29 token_chunk_offset_ptr,
30 + current_first_idx, # (batch,)
31 + current_last_idx, # (batch,)
32 + initial_state_idx, # (batch,)
33 + context_lens, # (batch,)
34 o_ptr, # (dim, seqlen) - actually pointing to x_ptr
35 # Matrix dimensions
30 - batch: tl.int32, # actually padded_batch
36 dim: tl.constexpr,
37 seqlen: tl.int32, # cu_seqlen
38 num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
39 # Strides
35 - stride_x_seq: tl.constexpr, # stride to get to next sequence,
40 stride_x_dim: tl.constexpr, # stride to get to next feature-value,
41 stride_x_token: tl.
42 constexpr, # stride to get to next token (same feature-index, same sequence-index)
@@ -42,18 +46,16 @@ def _causal_conv1d_fwd_kernel( # continuous batching
46 stride_istate_dim: tl.constexpr,
47 stride_istate_token: tl.constexpr,
48 stride_cache_indices: tl.constexpr,
45 - stride_o_seq: tl.constexpr,
49 stride_o_dim: tl.constexpr,
50 stride_o_token: tl.constexpr,
51 + stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M
52 # others
53 pad_slot_id: tl.constexpr,
54 # Meta-parameters
55 HAS_BIAS: tl.constexpr,
56 KERNEL_WIDTH: tl.constexpr,
57 SILU_ACTIVATION: tl.constexpr,
54 - HAS_INITIAL_STATES: tl.constexpr,
55 - HAS_CACHE: tl.constexpr,
56 - IS_CONTINUOUS_BATCHING: tl.constexpr,
58 + IS_APC_ENABLED: tl.constexpr,
59 USE_PAD_SLOT: tl.constexpr,
60 NP2_STATELEN: tl.constexpr,
61 BLOCK_M: tl.constexpr,
@@ -84,26 +86,57 @@ def _causal_conv1d_fwd_kernel( # continuous batching
86 # find the actual sequence length
87 seqlen = sequence_end_index - sequence_start_index
88
89 + B_size: tl.constexpr = (stride_block_m * BLOCK_M)
90 +
91 + if IS_APC_ENABLED:
92 + # Handle the case if prefix caching is enabled.
93 + # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr"
94 +
95 + # Get the length of the completed sequence so far and compute the offset.
96 + current_first_index = tl.load(current_first_idx + idx_seq)
97 + current_last_index = tl.load(current_last_idx + idx_seq)
98 + sequence_completed_index = tl.load(context_lens + idx_seq)
99 +
100 + # Compute the offset where the first stride_block_m-aligned first full block is
101 + # Value in "token-space"
102 + sequence_completed_offset_token = sequence_completed_index % B_size
103 + seq_completed_offset = B_size - sequence_completed_offset_token
104 + seq_end_offset = (seqlen - seq_completed_offset) % B_size
105 + last_full_block_token_index = sequence_end_index - seq_end_offset
106 + # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one
107 + if seq_end_offset == 0:
108 + last_full_block_token_index = last_full_block_token_index - B_size
109 +
110 + # Get the number of blocks to be filled for the current sequence
111 + # If n_block_to_fill = 0, then only the state at the sequence end is stored
112 + n_block_to_fill = current_last_index - current_first_index
113 +
114 + # Get the index of the init block
115 + conv_state_init_index = tl.load(initial_state_idx + idx_seq)
116 + else:
117 + n_block_to_fill = 0
118 + current_last_index = 0
119 + conv_state_init_index = 0
120 + current_first_index = 0
121 + last_full_block_token_index = 0
122 +
123 token_offset = BLOCK_M * chunk_offset
124 segment_len = min(BLOCK_M, seqlen - token_offset)
125
126 # base of the sequence
127 x_base = x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim # [BLOCK_N,]
128
93 - if IS_CONTINUOUS_BATCHING:
94 - # cache_idx
95 - conv_state_batch_coord = tl.load(conv_state_indices_ptr +
96 - idx_seq * stride_cache_indices).to(
97 - tl.int64)
98 - else:
99 - # cache_idx
100 - conv_state_batch_coord = idx_seq
129 + # cache_idx
130 + conv_states_input_coord = tl.load(conv_state_indices_ptr +
131 + idx_seq * stride_cache_indices +
132 + conv_state_init_index).to(tl.int64)
133 +
134 if USE_PAD_SLOT: # noqa
102 - if conv_state_batch_coord == pad_slot_id:
135 + if conv_states_input_coord == pad_slot_id:
136 # not processing as this is not the actual sequence
137 return
138 conv_states_base = (conv_states_ptr +
106 - (conv_state_batch_coord * stride_conv_state_seq) +
139 + (conv_states_input_coord * stride_conv_state_seq) +
140 (idx_feats * stride_conv_state_dim)) # [BLOCK_N,]
141
142 w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,]
@@ -113,10 +146,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
146 # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0]
147 if chunk_offset == 0:
148 # read from conv_states
116 - load_init_state = False
117 - if HAS_INITIAL_STATES: # the new HAS_INITIAL_STATES
118 - load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(
119 - tl.int1)
149 + load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1)
150 if load_init_state:
151 # load from conv_states
152 prior_tokens = conv_states_base + (state_len -
@@ -175,15 +205,23 @@ def _causal_conv1d_fwd_kernel( # continuous batching
205 (idx_feats < dim)[None, :]
206 ) # token-index # token-index # feature-index
207 loaded_x = tl.load(x_ptrs, mask_x, 0.0)
178 - new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
208 idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
180 - conv_states_ptrs_target = conv_states_base[None, :] + (
181 - idx_tokens_conv * stride_conv_state_tok)[:, None]
209 +
210 + # Compute the offset where the last block should be written in the conv_states
211 + conv_states_output_coord = tl.load(conv_state_indices_ptr +
212 + idx_seq * stride_cache_indices +
213 + current_last_index).to(tl.int64)
214 +
215 + conv_states_ptrs_target = (
216 + conv_states_ptr + (conv_states_output_coord *
217 + stride_conv_state_seq) + # Offset from seq
218 + (idx_feats * stride_conv_state_dim))[None, :] + ( # [BLOCK_N,]
219 + idx_tokens_conv * stride_conv_state_tok)[:, None]
220
221 mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats
222 < dim)[None, :]
223 tl.debug_barrier() # NOTE: use this due to bug in Triton compiler
186 - tl.store(conv_states_ptrs_target, new_conv_state, mask)
224 + tl.store(conv_states_ptrs_target, loaded_x, mask)
225
226 else:
227 if load_init_state:
@@ -192,12 +230,12 @@ def _causal_conv1d_fwd_kernel( # continuous batching
230
231 conv_states_ptrs_source = (
232 conv_states_ptr +
195 - (conv_state_batch_coord * stride_conv_state_seq) +
233 + (conv_states_input_coord * stride_conv_state_seq) +
234 (idx_feats * stride_conv_state_dim)[None, :] +
235 ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:,
236 None]
237 ) # [BLOCK_M, BLOCK_N]
200 - mask = ((conv_state_batch_coord < num_cache_lines)
238 + mask = ((conv_states_input_coord < num_cache_lines)
239 & ((idx_tokens_conv + seqlen) < state_len)[:, None]
240 & (idx_feats < dim)[None, :])
241 conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
@@ -280,6 +318,45 @@ def _causal_conv1d_fwd_kernel( # continuous batching
318 conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N]
319 col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
320
321 + # Store intermediate states aligned with stride_block_m
322 + # The additional states are cached starting from the last stride_block_m.
323 + # For example:
324 + # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved.
325 + # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last
326 + # stride_block_m are cached.
327 + # For example chunk_offset = n_block_to_fill stores the state at last_full_block
328 + if (chunk_offset - 1) < n_block_to_fill:
329 + # Store the states at the chunk boundaries from the start of the sequence
330 + idx_tokens_last = (last_full_block_token_index -
331 + (n_block_to_fill - chunk_offset) * B_size -
332 + state_len) + tl.arange(
333 + 0, NP2_STATELEN) # [BLOCK_M]
334 + x_ptrs = x_ptr + (idx_tokens_last * stride_x_token)[:, None] + (
335 + idx_feats * stride_x_dim)[None, :] # [BLOCK_M,BLOCK_N,]
336 +
337 + mask_x = (
338 + (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[None, :]
339 + ) # token-index # token-index # feature-index
340 + loaded_x = tl.load(x_ptrs, mask_x, 0.0)
341 + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
342 +
343 + # cache_idx
344 + conv_states_output_coord = tl.load(conv_state_indices_ptr +
345 + idx_seq * stride_cache_indices +
346 + current_first_index +
347 + (chunk_offset - 1)).to(tl.int64)
348 +
349 + conv_states_ptrs_target = (
350 + conv_states_ptr + (conv_states_output_coord *
351 + stride_conv_state_seq) + # Offset from seq
352 + (idx_feats * stride_conv_state_dim))[None, :] + ( # [BLOCK_N,]
353 + idx_tokens_conv * stride_conv_state_tok)[:, None]
354 +
355 + mask = (idx_tokens_conv < state_len)[:, None] & \
356 + (idx_feats < dim)[None, :]
357 + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler
358 + tl.store(conv_states_ptrs_target, loaded_x, mask)
359 +
360 if HAS_BIAS:
361 bias = bias_ptr + idx_feats
362 mask_bias = idx_feats < dim
@@ -368,6 +445,11 @@ def causal_conv1d_fn(
445 has_initial_state: Optional[torch.Tensor] = None,
446 activation: Optional[str] = "silu",
447 pad_slot_id: int = PAD_SLOT_ID,
448 + current_first_idx: Optional[torch.Tensor] = None,
449 + current_last_idx: Optional[torch.Tensor] = None,
450 + initial_state_idx: Optional[torch.Tensor] = None,
451 + context_lens: Optional[torch.Tensor] = None,
452 + block_size_to_align=0,
453 metadata=None,
454 validate_data=False,
455 ):
@@ -378,7 +460,7 @@ def causal_conv1d_fn(
460 sequences are concatenated from left to right for varlen
461 weight: (dim, width)
462 conv_states: (...,dim,width - 1) itype
381 - updated inplace if provided
463 + updated inplace if cache_indices are not provided
464 [it use `cache_indices` to get the index to the cache of conv_state for that sequence
465
466 conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True
@@ -410,7 +492,16 @@ def causal_conv1d_fn(
492 for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
493 in this case, the kernel will not process entries at
494 indices 0 and 3
413 -
495 + current_first_idx: (batch,), dtype int32
496 + The pointer into cache_indices, where the first cache block to be filled is located.
497 + current_last_idx: (batch,), dtype int32
498 + The pointer into cache_indices, where the last cache block to be filled is located.
499 + initial_state_idx: (batch,), dtype int32
500 + The pointer into cache_indices, where the cache block containing the initial state is located.
501 + context_lens: (batch,), dtype int32
502 + The number of tokens already completed for each sequence
503 + block_size_to_align: int
504 + The block size to align the cached states to
505 out: same shape as `x`
506 """
507 if isinstance(activation, bool) and activation:
@@ -451,7 +542,6 @@ def causal_conv1d_fn(
542 np2_statelen = triton.next_power_of_2(state_len)
543
544 padded_batch = query_start_loc.size(0) - 1
454 - stride_x_seq = 0
545 stride_x_dim = x.stride(0)
546 stride_x_token = x.stride(1)
547 stride_w_dim = weight.stride(0)
@@ -460,6 +550,7 @@ def causal_conv1d_fn(
550 stride_istate_dim = 0
551 stride_istate_token = 0
552 num_cache_lines = 0
553 + BLOCK_M = 8
554 if conv_states is not None:
555 # extensions to support vLLM:
556 # 1. conv_states is used to replaced initial_states
@@ -475,11 +566,9 @@ def causal_conv1d_fn(
566 stride_istate_token = conv_states.stride(2)
567 assert stride_istate_dim == 1
568 if out.dim() == 2:
478 - stride_o_seq = 0
569 stride_o_dim = out.stride(0)
570 stride_o_token = out.stride(1)
571 else:
482 - stride_o_seq = out.stride(0)
572 stride_o_dim = out.stride(1)
573 stride_o_token = out.stride(2)
574 stride_cache_indices = cache_indices.stride(
@@ -502,6 +591,12 @@ def causal_conv1d_fn(
591 assert weight.stride(1) == 1
592 assert (dim, width) == weight.shape
593 assert is_channel_last, "Need to run in channel-last layout"
594 + if block_size_to_align is not None and block_size_to_align > 0:
595 + assert (
596 + block_size_to_align % BLOCK_M
597 + ) == 0, "The mamba block size needs to be divisible by the BLOCK_M"
598 + else:
599 + block_size_to_align = BLOCK_M
600
601 if metadata is None:
602
@@ -584,14 +679,16 @@ def causal_conv1d_fn(
679 query_start_loc,
680 batch_ptr,
681 token_chunk_offset_ptr,
682 + current_first_idx,
683 + current_last_idx,
684 + initial_state_idx,
685 + context_lens,
686 out,
687 # Matrix dimensions
589 - padded_batch,
688 dim,
689 cu_seqlen,
690 num_cache_lines,
691 # stride
594 - stride_x_seq,
692 stride_x_dim,
693 stride_x_token,
694 stride_w_dim,
@@ -600,22 +697,20 @@ def causal_conv1d_fn(
697 stride_istate_dim,
698 stride_istate_token,
699 stride_cache_indices,
603 - stride_o_seq,
700 stride_o_dim,
701 stride_o_token,
702 + block_size_to_align // BLOCK_M,
703 # others
704 pad_slot_id,
705 # META
706 HAS_BIAS=bias is not None,
707 KERNEL_WIDTH=width,
708 SILU_ACTIVATION=activation in ["silu", "swish"],
612 - HAS_INITIAL_STATES=has_initial_state is not None,
613 - HAS_CACHE=conv_states is not None,
614 - IS_CONTINUOUS_BATCHING=cache_indices is not None,
709 + IS_APC_ENABLED=current_last_idx is not None,
710 USE_PAD_SLOT=pad_slot_id is not None,
711 NP2_STATELEN=np2_statelen,
712 #launch_cooperative_grid=True
618 - BLOCK_M=8,
713 + BLOCK_M=BLOCK_M,
714 BLOCK_N=256,
715 num_stages=2,
716 )
@@ -629,10 +724,11 @@ def _causal_conv1d_update_kernel(
724 w_ptr, # (dim, width)
725 bias_ptr,
726 conv_state_ptr,
632 - cache_seqlens_ptr, # circular buffer
727 conv_state_indices_ptr,
728 num_accepted_tokens_ptr,
729 query_start_loc_ptr, # (batch + 1)
730 + current_last_idx, # (batch,)
731 + initial_state_idx, #(batch,)
732 o_ptr, # (batch, dim, seqlen)
733 # Matrix dimensions
734 batch: int,
@@ -660,7 +756,7 @@ def _causal_conv1d_update_kernel(
756 KERNEL_WIDTH: tl.constexpr,
757 SILU_ACTIVATION: tl.constexpr,
758 IS_VARLEN: tl.constexpr,
663 - IS_CONTINUOUS_BATCHING: tl.constexpr,
759 + IS_APC_ENABLED: tl.constexpr,
760 IS_SPEC_DECODING: tl.constexpr,
761 NP2_STATELEN: tl.constexpr,
762 USE_PAD_SLOT: tl.constexpr,
@@ -674,15 +770,21 @@ def _causal_conv1d_update_kernel(
770 # [BLOCK_N,] elements along the feature-dimension (channel)
771 idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
772
677 - if IS_CONTINUOUS_BATCHING:
678 - # mask = idx_seq < batch
679 - conv_state_batch_coord = tl.load(conv_state_indices_ptr +
680 - idx_seq * stride_state_indices).to(
681 - tl.int64)
773 + if IS_APC_ENABLED:
774 + # Get the state from the initial_state_idx
775 + conv_state_init = tl.load(initial_state_idx + idx_seq)
776 + current_last_index = tl.load(current_last_idx + idx_seq)
777 else:
683 - conv_state_batch_coord = idx_seq
778 + conv_state_init = 0
779 + current_last_index = 0
780 +
781 + # cache_idx
782 + conv_states_input_coord = tl.load(conv_state_indices_ptr +
783 + idx_seq * stride_state_indices +
784 + conv_state_init).to(tl.int64)
785 +
786 if USE_PAD_SLOT: # noqa
685 - if conv_state_batch_coord == pad_slot_id:
787 + if conv_states_input_coord == pad_slot_id:
788 # not processing as this is not the actual sequence
789 return
790
@@ -726,7 +828,7 @@ def _causal_conv1d_update_kernel(
828
829 # STEP 1: READ init_state data
830 conv_states_base = (conv_state_ptr +
729 - (conv_state_batch_coord * stride_conv_state_seq) +
831 + (conv_states_input_coord * stride_conv_state_seq) +
832 (idx_feats * stride_conv_state_dim))
833 mask_w = idx_feats < dim
834
@@ -754,12 +856,12 @@ def _causal_conv1d_update_kernel(
856 # window manner, at each forward pass, the tokens are shift by 1, so we
857 # load since idx_tokens + 1.
858 conv_state_ptrs_source = (
757 - conv_state_ptr + (conv_state_batch_coord * stride_conv_state_seq) +
859 + conv_state_ptr + (conv_states_input_coord * stride_conv_state_seq) +
860 conv_state_token_offset * stride_conv_state_tok +
861 (idx_feats * stride_conv_state_dim)[None, :] +
862 ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) *
863 stride_conv_state_tok)[:, None]) # [BLOCK_M, BLOCK_N]
762 - mask = ((conv_state_batch_coord < num_cache_lines)
864 + mask = ((conv_states_input_coord < num_cache_lines)
865 & ((idx_tokens + seqlen) < state_len)[:, None]
866 & (idx_feats < dim)[None, :])
867 conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
@@ -778,11 +880,16 @@ def _causal_conv1d_update_kernel(
880
881 new_conv_state = tl.where(mask, conv_state, loaded_x)
882
781 - conv_state_base = (conv_state_ptr +
782 - (conv_state_batch_coord * stride_conv_state_seq) +
783 - (idx_feats * stride_conv_state_dim)) # [BLOCK_N,]
784 - conv_state_ptrs_target = conv_state_base + (
785 - idx_tokens * stride_conv_state_tok)[:, None] # [BLOCK_M, BLOCK_N]
883 + # Get the state from the initial_state_idx
884 + # cache_idx
885 + conv_states_offset = tl.load(conv_state_indices_ptr +
886 + idx_seq * stride_state_indices +
887 + current_last_index).to(tl.int64)
888 + conv_state_ptrs_target = (
889 + conv_state_ptr +
890 + (conv_states_offset * stride_conv_state_seq) + # Offset from seq
891 + (idx_feats * stride_conv_state_dim))[None, :] + ( # [BLOCK_N,]
892 + idx_tokens * stride_conv_state_tok)[:, None]
893 mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
894 tl.store(conv_state_ptrs_target, new_conv_state, mask)
895
@@ -923,12 +1030,13 @@ def causal_conv1d_update(
1030 weight: torch.Tensor,
1031 bias: Optional[torch.Tensor] = None,
1032 activation: Union[bool, str, None] = None,
926 - cache_seqlens: Optional[torch.Tensor] = None,
1033 conv_state_indices: Optional[torch.Tensor] = None,
1034 num_accepted_tokens: Optional[torch.Tensor] = None,
1035 query_start_loc: Optional[torch.Tensor] = None,
1036 max_query_len: int = -1,
1037 pad_slot_id: int = PAD_SLOT_ID,
1038 + current_last_idx: Optional[torch.Tensor] = None,
1039 + initial_state_idx: Optional[torch.Tensor] = None,
1040 validate_data=False,
1041 ):
1042 """
@@ -942,15 +1050,14 @@ def causal_conv1d_update(
1050 conv_state: (..., dim, state_len), where state_len >= width - 1
1051 weight: (dim, width)
1052 bias: (dim,)
945 - cache_seqlens: (batch,), dtype int32.
946 - If not None, the conv_state is treated as a circular buffer.
947 - The conv_state will be updated by copying x to the conv_state
948 - starting at the index
949 - @cache_seqlens % state_len.
1053 conv_state_indices: (batch,), dtype int32
1054 If not None, the conv_state is a larger tensor along the batch dim,
1055 and we are selecting the batch coords specified by conv_state_indices.
1056 Useful for a continuous batching scenario.
1057 + current_last_idx: (batch,), dtype int32
1058 + The pointer into conv_state_indices, where the last cache block to be filled is located.
1059 + initial_state_idx: (batch,), dtype int32
1060 + The pointer into conv_state_indices, where the cache block containing the initial state is located.
1061 num_accepted_tokens: (batch,), dtype int32
1062 If not None, it indicates the number of accepted tokens for each
1063 sequence in the batch.
@@ -963,15 +1070,14 @@ def causal_conv1d_update(
1070 If query_start_loc is not None, this indicates the maximum query
1071 length in the batch.
1072 pad_slot_id: int
966 - if cache_indices is passed, lets the kernel identify padded
1073 + if conv_state_indices is passed, lets the kernel identify padded
1074 entries that will not be processed,
968 - for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
1075 + for example: conv_state_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
1076 in this case, the kernel will not process entries at
1077 indices 0 and 3
1078 out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x`
1079 """
1080 if validate_data:
974 - assert cache_seqlens is None # not implemented yet - ok for vLLM
1081 assert pad_slot_id is not None
1082 assert x.stride(1) == 1
1083 if isinstance(activation, bool):
@@ -1011,7 +1117,6 @@ def causal_conv1d_update(
1117
1118 assert num_cache_lines >= batch
1119 assert weight.stride(1) == 1 # Need this
1014 - assert cache_seqlens is None # not needed for vLLM - circular buffer
1120
1121 # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o'
1122 out = x
@@ -1050,10 +1155,11 @@ def causal_conv1d_update(
1155 weight,
1156 bias,
1157 conv_state,
1053 - cache_seqlens,
1158 conv_state_indices,
1159 num_accepted_tokens,
1160 query_start_loc,
1161 + current_last_idx,
1162 + initial_state_idx,
1163 out,
1164 # Matrix dimensions
1165 batch,
@@ -1081,7 +1187,7 @@ def causal_conv1d_update(
1187 KERNEL_WIDTH=width,
1188 SILU_ACTIVATION=activation in ["silu", "swish"],
1189 IS_VARLEN=query_start_loc is not None,
1084 - IS_CONTINUOUS_BATCHING=conv_state_indices is not None,
1190 + IS_APC_ENABLED=current_last_idx is not None,
1191 IS_SPEC_DECODING=num_accepted_tokens is not None,
1192 NP2_STATELEN=np2_statelen,
1193 USE_PAD_SLOT=pad_slot_id is not None,
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+17 -1
@@ -52,6 +52,7 @@ def _selective_scan_update_kernel(
52 z_ptr,
53 out_ptr,
54 state_batch_indices_ptr,
55 + dst_state_batch_indices_ptr,
56 pad_slot_id,
57 # Matrix dimensions
58 batch,
@@ -107,11 +108,17 @@ def _selective_scan_update_kernel(
108 # is taken from the state_batch_indices_ptr Otherwise, the state coordinate
109 # is the same as the batch id.
110 if HAS_STATE_BATCH_INDICES:
111 + dst_state_batch_indices_ptr += pid_b
112 + dst_state_batch_idx = tl.load(dst_state_batch_indices_ptr).to(tl.int64)
113 + dst_state_ptr = state_ptr + (dst_state_batch_idx * stride_state_batch +
114 + pid_h * stride_state_head)
115 state_batch_indices_ptr += pid_b
116 state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64)
117 state_ptr += (state_batch_idx * stride_state_batch +
118 pid_h * stride_state_head)
119 else:
120 + dst_state_ptr = state_ptr + pid_b * stride_state_batch + \
121 + pid_h * stride_state_head
122 state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head
123
124 x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head
@@ -131,6 +138,8 @@ def _selective_scan_update_kernel(
138 offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)
139 state_ptrs = state_ptr + (offs_m[:, None] * stride_state_dim +
140 offs_n[None, :] * stride_state_dstate)
141 + dst_state_ptrs = dst_state_ptr + (offs_m[:, None] * stride_state_dim +
142 + offs_n[None, :] * stride_state_dstate)
143 x_ptrs = x_ptr + offs_m * stride_x_dim
144 dt_ptrs = dt_ptr + offs_m * stride_dt_dim
145 if HAS_DT_BIAS:
@@ -185,7 +194,7 @@ def _selective_scan_update_kernel(
194 mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)
195 if HAS_STATE_BATCH_INDICES:
196 mask &= (state_batch_idx != pad_slot_id)
188 - tl.store(state_ptrs, state, mask=mask)
197 + tl.store(dst_state_ptrs, state, mask=mask)
198 out = tl.sum(state * C[None, :], axis=1)
199 if HAS_D:
200 out += x * D
@@ -205,6 +214,7 @@ def selective_state_update(state,
214 dt_bias=None,
215 dt_softplus=False,
216 state_batch_indices=None,
217 + dst_state_batch_indices=None,
218 pad_slot_id=PAD_SLOT_ID,
219 out=None):
220 """
@@ -266,6 +276,11 @@ def selective_state_update(state,
276 assert dt_bias.shape == (nheads, dim)
277 if state_batch_indices is not None:
278 assert state_batch_indices.shape == (batch, )
279 + if dst_state_batch_indices is not None:
280 + assert dst_state_batch_indices.shape == (batch, )
281 + else:
282 + # revert to the default behavior of in-place state updates
283 + dst_state_batch_indices = state_batch_indices
284 assert out.shape == x.shape
285
286 grid = lambda META: (triton.cdiv(dim, META['BLOCK_SIZE_M']), batch, nheads)
@@ -292,6 +307,7 @@ def selective_state_update(state,
307 z,
308 out,
309 state_batch_indices,
310 + dst_state_batch_indices,
311 pad_slot_id,
312 batch,
313 nheads,
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+25 -19
@@ -35,6 +35,7 @@ def _mamba_chunk_scan_combined_fwd(x,
35 z=None,
36 dt_bias=None,
37 initial_states=None,
38 + return_intermediate_states=False,
39 seq_idx=None,
40 cu_seqlens=None,
41 cu_chunk_seqlens=None,
@@ -151,28 +152,32 @@ def _mamba_chunk_scan_combined_fwd(x,
152 initial_states=initial_states,
153 )
154
154 - return states[last_chunk_indices]
155 + if return_intermediate_states:
156 + return states
157 + else:
158 + return states[last_chunk_indices]
159
160
161 def mamba_chunk_scan_combined_varlen(
158 - x,
159 - dt,
160 - A,
161 - B,
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,
173 - dt_softplus=False,
174 - dt_limit=(0.0, float("inf")),
175 - state_dtype=None,
162 + x,
163 + dt,
164 + A,
165 + B,
166 + C,
167 + chunk_size,
168 + cu_seqlens,
169 + cu_chunk_seqlens,
170 + last_chunk_indices,
171 + seq_idx,
172 + out,
173 + D=None,
174 + z=None,
175 + dt_bias=None,
176 + initial_states=None,
177 + dt_softplus=False,
178 + dt_limit=(0.0, float("inf")),
179 + return_intermediate_states=False,
180 + state_dtype=None,
181 ):
182 """
183 Argument:
@@ -213,6 +218,7 @@ def mamba_chunk_scan_combined_varlen(
218 z=z,
219 dt_bias=dt_bias,
220 initial_states=initial_states,
221 + return_intermediate_states=return_intermediate_states,
222 seq_idx=seq_idx,
223 cu_seqlens=cu_seqlens,
224 cu_chunk_seqlens=cu_chunk_seqlens,
vllm/model_executor/models/bamba.py
-4
@@ -453,12 +453,8 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
453 config = vllm_config.model_config.hf_config
454 self.vllm_config = vllm_config
455 self.model_config = vllm_config.model_config
456 - cache_config = vllm_config.cache_config
456 lora_config = vllm_config.lora_config
457 scheduler_config = vllm_config.scheduler_config
459 - assert not cache_config.enable_prefix_caching, \
460 - "Bamba currently does not support prefix caching"
461 -
458 self.quant_config = vllm_config.quant_config
459
460 super().__init__()
vllm/model_executor/models/config.py
+59 -10
@@ -292,10 +292,33 @@ class MambaModelConfig(VerifyAndUpdateConfig):
292 cache_config = vllm_config.cache_config
293 compilation_config = vllm_config.compilation_config
294
295 - # TODO(tdoublep): remove once prefix caching is enabled
296 - cache_config.enable_prefix_caching = False
297 - logger.info("Hybrid or mamba-based model detected: disabling prefix "
298 - "caching since it is not yet supported.")
295 + # Set mamba block size to max_model_len (this may get
296 + # override by prefix caching logic later)
297 + cache_config.mamba_block_size = model_config.max_model_len
298 +
299 + # TODO(@tdoublep) find a better way to do this than whitelist
300 + MAMBA2_MODELS = [
301 + "BambaForCausalLM",
302 + "FalconH1ForCausalLM",
303 + "GraniteMoeHybridForCausalLM",
304 + "Mamba2ForCausalLM",
305 + "NemotronHForCausalLM",
306 + "Zamba2ForCausalLM",
307 + ]
308 + if cache_config.enable_prefix_caching:
309 + if model_config.architecture in MAMBA2_MODELS:
310 + logger.info("Warning: Prefix caching is currently enabled. "
311 + "Its support for Mamba2 layers is experimental. "
312 + "Please report any issues you may observe.")
313 + else:
314 + logger.info("Hybrid or mamba-based model detected without "
315 + "support for prefix caching: disabling.")
316 + cache_config.enable_prefix_caching = False
317 +
318 + # TODO(tdoublep): remove once cascade attention is supported
319 + logger.info("Disabling cascade attention since it is not supported "
320 + "for hybrid models.")
321 + model_config.disable_cascade_attn = True
322
323 # TODO(tdoublep): remove as full cuda graph support is added
324 FCG_NOT_SUPPORTED_MODELS = [
@@ -360,12 +383,38 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
383 block_size=model_config.max_model_len,
384 ).page_size_bytes
385
363 - # some attention backends (e.g. FA) only support setting
364 - # block size to multiple of 16, so let's suggest a value
365 - # that would work (note: FA is currently not compatible
366 - # with mamba layers, use FlashInfer instead).
367 - attn_block_size = 16 * cdiv(mamba_page_size,
368 - 16 * attn_page_size_1_token)
386 + if cache_config.enable_prefix_caching:
387 + # With prefix caching, select attention block size to
388 + # optimize for mamba kernel performance
389 +
390 + # mamba SSD kernel uses a chunk_size, e.g. 256
391 + # Align the block to the kernel: use lowest multiple of chunk_size
392 + # of attention tokens that would fit mamba_page_size:
393 + # e.g. for mamba page size = 788kB
394 + # attn_1_token = 2kB -> fits ~394 tokens
395 + # then round up to a mulitple of 256 -> 512 tokens
396 + # End result:
397 + # attn_block_size = 512
398 + # mamba_block_size = 512 (aligned to a multiple of chunk_size)
399 + # TODO(tdoublep): this constraint can be relaxed fairly
400 + # easily by changing the way we layout chunks in the
401 + # mamba2 kernels.
402 + chunk_size = model_config.get_mamba_chunk_size()
403 + attn_tokens_per_mamba_state = \
404 + cdiv(mamba_page_size, attn_page_size_1_token)
405 + attn_block_size = chunk_size * \
406 + cdiv(attn_tokens_per_mamba_state, chunk_size)
407 + cache_config.mamba_block_size = attn_block_size
408 + else:
409 + # Without prefix caching, select minimum valid attention block size
410 + # to minimize mamba state padding
411 +
412 + # some attention backends (e.g. FA) only support setting
413 + # block size to multiple of 16, so let's suggest a value
414 + # that would work (note: FA is currently not compatible
415 + # with mamba layers, use FlashInfer instead).
416 + attn_block_size = 16 * cdiv(mamba_page_size,
417 + 16 * attn_page_size_1_token)
418
419 # override attention block size if either (a) the
420 # user has not set it or (b) the user has set it
vllm/model_executor/models/falcon_h1.py
-3
@@ -540,11 +540,8 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
540 config = vllm_config.model_config.hf_config
541 self.vllm_config = vllm_config
542 self.model_config = vllm_config.model_config
543 - cache_config = vllm_config.cache_config
543 lora_config = vllm_config.lora_config
544 scheduler_config = vllm_config.scheduler_config
546 - assert (not cache_config.enable_prefix_caching
547 - ), "FalconH1 currently does not support prefix caching"
545
546 self.quant_config = vllm_config.quant_config
547
vllm/model_executor/models/granitemoehybrid.py
-5
@@ -549,13 +549,8 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
549 config = vllm_config.model_config.hf_config
550 self.vllm_config = vllm_config
551 self.model_config = vllm_config.model_config
552 - cache_config = vllm_config.cache_config
552 lora_config = vllm_config.lora_config
553 scheduler_config = vllm_config.scheduler_config
555 - if cache_config.enable_prefix_caching:
556 - raise RuntimeError(
557 - "GraniteMoeHybrid currently does not support prefix caching")
558 -
554 self.quant_config = vllm_config.quant_config
555 self.config = config
556 self.scheduler_config = scheduler_config
vllm/model_executor/models/mamba2.py
-3
@@ -222,11 +222,8 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
222
223 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
224 config = vllm_config.model_config.hf_config
225 - cache_config = vllm_config.cache_config
225 lora_config = vllm_config.lora_config
226 scheduler_config = vllm_config.scheduler_config
228 - assert not cache_config.enable_prefix_caching, \
229 - "Mamba does not support prefix caching"
227
228 super().__init__()
229 self.config = config
vllm/model_executor/models/nemotron_h.py
-3
@@ -505,11 +505,8 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
505 config = vllm_config.model_config.hf_config
506 self.vllm_config = vllm_config
507 self.model_config = vllm_config.model_config
508 - cache_config = vllm_config.cache_config
508 lora_config = vllm_config.lora_config
509 scheduler_config = vllm_config.scheduler_config
511 - assert not cache_config.enable_prefix_caching, \
512 - "NemotronH currently does not support prefix caching"
510
511 self.quant_config = vllm_config.quant_config
512
vllm/model_executor/models/zamba2.py
-3
@@ -868,11 +868,8 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
868 (not supported by Mamba)
869 """
870 config = vllm_config.model_config.hf_config
871 - cache_config = vllm_config.cache_config
871 lora_config = vllm_config.lora_config
872 scheduler_config = vllm_config.scheduler_config
874 - assert not cache_config.enable_prefix_caching, \
875 - "Mamba does not support prefix caching"
873
874 super().__init__()
875 self.config = config
vllm/v1/attention/backends/mamba2_attn.py
+90 -1
@@ -122,6 +122,11 @@ class Mamba2AttentionMetadata:
122 last_chunk_indices_p: Optional[torch.Tensor]
123
124 state_indices_tensor: torch.Tensor # shape: [batch,]
125 + current_last_idx: torch.Tensor
126 + current_first_idx_p: torch.Tensor
127 + last_state_idx: torch.Tensor
128 + context_lens_p: torch.Tensor
129 + last_computed_offset_p: torch.Tensor
130
131 # The following attributes are for triton implementation of causal_conv1d
132 nums_dict: Optional[dict] = None
@@ -138,6 +143,24 @@ class Mamba2AttentionMetadataBuilder(
143 self.chunk_size = vllm_config.model_config.get_mamba_chunk_size()
144 assert self.chunk_size is not None, (
145 "chunk_size needs to be set in the model config for Mamba2 models")
146 + if self.vllm_config.cache_config.enable_prefix_caching:
147 + self.state_indices_tensor = torch.empty(
148 + (self.decode_cudagraph_max_bs,
149 + cdiv(vllm_config.model_config.max_model_len,
150 + kv_cache_spec.block_size)),
151 + dtype=torch.int32,
152 + device=device,
153 + )
154 + self.current_last_idx = torch.empty(
155 + (self.decode_cudagraph_max_bs, ),
156 + dtype=torch.int32,
157 + device=device,
158 + )
159 + self.last_state_idx = torch.empty(
160 + (self.decode_cudagraph_max_bs, ),
161 + dtype=torch.int32,
162 + device=device,
163 + )
164
165 def build(self,
166 common_prefix_len: int,
@@ -158,7 +181,45 @@ class Mamba2AttentionMetadataBuilder(
181 # for causal_conv1d
182 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
183
161 - state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
184 + context_lens, context_lens_p = None, None
185 + current_first_idx, current_first_idx_p = None, None
186 + last_computed_offset, last_computed_offset_p = None, None
187 +
188 + if self.vllm_config.cache_config.enable_prefix_caching:
189 + # Return a tensor of shape (#requests, #max blocks)
190 + state_indices_tensor = common_attn_metadata.block_table_tensor
191 +
192 + # Additional cache-related varaiables:
193 + mamba_block_size = self.kv_cache_spec.block_size
194 + seq_lens_pending = (
195 + torch.roll(common_attn_metadata.query_start_loc, -1, -1) -
196 + common_attn_metadata.query_start_loc)[:-1]
197 + context_lens = common_attn_metadata.seq_lens - \
198 + seq_lens_pending
199 + last_computed_offset = \
200 + context_lens % mamba_block_size
201 + # Indices: last_computed <= current_first <= current_last
202 + # Cases:
203 + # last_computed == current_first if last state was partially
204 + # computed and needs to be updated
205 + # current_first == current_last if no block crossing occurs, and
206 + # only one state will be stored
207 + # 0th based indexing leads to "-1" -> e.g. 16 computed -> state[15]:
208 + current_last_idx = cdiv(context_lens + seq_lens_pending,
209 + mamba_block_size) - 1
210 + current_first_idx = cdiv(context_lens + 1, mamba_block_size) - 1
211 + last_state_idx = cdiv(context_lens, mamba_block_size) - 1
212 + # -1 in case it's non-computed and causes later issues with indexing
213 + last_state_idx = \
214 + last_state_idx.clamp(min=0)
215 +
216 + else:
217 + # Always return just a single block per each request:
218 + state_indices_tensor = common_attn_metadata.block_table_tensor[:,
219 + 0]
220 + # Additional cache-related varaiables:
221 + current_last_idx = None
222 + last_state_idx = None
223
224 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
225 split_decodes_and_prefills(
@@ -178,6 +239,16 @@ class Mamba2AttentionMetadataBuilder(
239 query_start_loc_p = common_attn_metadata.query_start_loc[
240 -num_prefills - 1:] - num_decode_tokens
241
242 + if self.vllm_config.cache_config.enable_prefix_caching:
243 + assert context_lens is not None
244 + context_lens_p = context_lens[num_reqs - num_prefills:num_reqs]
245 + assert last_computed_offset is not None
246 + last_computed_offset_p = last_computed_offset[
247 + num_reqs - num_prefills:num_reqs]
248 + assert current_first_idx is not None
249 + current_first_idx_p = current_first_idx[num_reqs -
250 + num_prefills:num_reqs]
251 +
252 num_computed_tokens_p = \
253 common_attn_metadata.num_computed_tokens_cpu[
254 num_reqs - num_prefills:num_reqs]
@@ -252,6 +323,19 @@ class Mamba2AttentionMetadataBuilder(
323 state_indices_tensor = self.state_indices_tensor[:num_input_tokens]
324 state_indices_tensor[num_decodes:] = PAD_SLOT_ID
325
326 + if self.vllm_config.cache_config.enable_prefix_caching:
327 + self.current_last_idx[:num_decodes].copy_(current_last_idx,
328 + non_blocking=True)
329 + current_last_idx = \
330 + self.current_last_idx[:num_input_tokens]
331 + current_last_idx[num_decodes:] = 0
332 +
333 + self.last_state_idx[:num_decodes].copy_(last_state_idx,
334 + non_blocking=True)
335 + last_state_idx = \
336 + self.last_state_idx[:num_input_tokens]
337 + last_state_idx[num_decodes:] = 0
338 +
339 attn_metadata = Mamba2AttentionMetadata(
340 num_prefills=num_prefills,
341 num_prefill_tokens=num_prefill_tokens,
@@ -269,5 +353,10 @@ class Mamba2AttentionMetadataBuilder(
353 nums_dict=nums_dict,
354 batch_ptr=batch_ptr,
355 token_chunk_offset_ptr=token_chunk_offset_ptr,
356 + current_last_idx=current_last_idx,
357 + current_first_idx_p=current_first_idx_p,
358 + last_state_idx=last_state_idx,
359 + context_lens_p=context_lens_p,
360 + last_computed_offset_p=last_computed_offset_p,
361 )
362 return attn_metadata
vllm/v1/core/single_type_kv_cache_manager.py
+22 -4
@@ -546,20 +546,38 @@ class MambaManager(SingleTypeKVCacheManager):
546 kv_cache_spec,
547 MambaSpec), ("MambaManager can only be used for mamba groups")
548 assert dcp_world_size == 1, "DCP not support mamba now."
549 - # Prefix caching is not supported for mamba now. Always return empty
550 - # list.
549 computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
550 [] for _ in range(len(kv_cache_group_ids)))
551 +
552 + max_num_blocks = max_length // kv_cache_spec.block_size
553 + # Search from right to left and early stop when a match is found.
554 + for i in range(max_num_blocks - 1, -1, -1):
555 + if cached_block := block_pool.get_cached_block(
556 + block_hashes[i], kv_cache_group_ids):
557 + for computed, cached in zip(computed_blocks, cached_block):
558 + # the hit length logic later assumes:
559 + # hit_length = len(hit_blocks_other_attn[0])
560 + # * self.other_block_size
561 + # so we insert dummy blocks at the beginning:
562 + if i > 0:
563 + computed.extend([block_pool.null_block] * i)
564 + computed.append(cached)
565 + break # we just need the last match - early stopping
566 +
567 return computed_blocks
568
569 def remove_skipped_blocks(self, request_id: str,
570 num_computed_tokens: int) -> None:
557 - # Each request will always have 1 block at this moment, so no need to
558 - # remove blocks.
571 + # Here unused blocks may be freed up for running requests.
572 + # TODO(@s3woz) Free up all blocks that aren't needed by Mamba2
573 + # (for which find_longest_cache_hit returns block_pool.null_block)
574 pass
575
576 def get_num_common_prefix_blocks(self, request_id: str,
577 num_running_requests: int) -> int:
578 + """
579 + cascade attention is not supported by mamba
580 + """
581 return 0
582
583 def get_num_blocks_to_allocate(
vllm/v1/kv_cache_interface.py
+2 -4
@@ -233,10 +233,8 @@ class MambaSpec(KVCacheSpec):
233 return page_size
234
235 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
236 - # We allocate 1 block for each request now, so max_memory_usage_bytes is
237 - # the same as page_size_bytes.
238 - # Need to update this when supporting prefix caching.
239 - return self.page_size_bytes
236 + max_model_len = vllm_config.model_config.max_model_len
237 + return cdiv(max_model_len, self.block_size) * self.page_size_bytes
238
239
240 @dataclass(frozen=True)
vllm/v1/worker/gpu_model_runner.py
+2 -8
@@ -4240,21 +4240,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4240 not in ["qwen3_next"]):
4241 raise NotImplementedError(
4242 "Mamba with speculative decoding is not supported yet.")
4243 - if self.vllm_config.cache_config.enable_prefix_caching:
4244 - raise NotImplementedError(
4245 - "Prefix caching is not supported for Mamba yet.")
4246 - max_model_len = self.vllm_config.model_config.max_model_len
4247 -
4243 + mamba_block_size = self.vllm_config.cache_config.mamba_block_size
4244 page_size_padded = (
4245 self.vllm_config.cache_config.mamba_page_size_padded)
4246
4251 - # Set block_size to max_model_len, so that mamba model will always
4252 - # have only one block in the KV cache.
4247 for layer_name, mamba_module in mamba_layers.items():
4248 kv_cache_spec[layer_name] = MambaSpec(
4249 shapes=mamba_module.get_state_shape(),
4250 dtypes=mamba_module.get_state_dtype(),
4257 - block_size=max_model_len,
4251 + block_size=mamba_block_size,
4252 page_size_padded=page_size_padded,
4253 mamba_type=mamba_module.mamba_type,
4254 num_speculative_blocks=(