[XPU][CI] Make `test_mamba_prefix_cache` block-size agnostic (#58269)

Signed-off-by: Lin, Fanli <fanli.lin@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

Fanli Lin committed Sep 29, 2026 at 07:12 UTC 6a6c2952faecdce6760dafb30df1eff95c79d717
2 files changed +110 -103
tests/v1/e2e/general/test_mamba_prefix_cache.py
+103 -99
@@ -55,7 +55,10 @@ async_scheduling_mode = False
55 num_accepted_tokens = 1
56 prompt_token_ids: list[int] = []
57 MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
58 -BLOCK_SIZE = 560
58 +# Must be a value no platform rewrites, since the step actions below encode
59 +# where block boundaries fall. XPU rounds GDN block sizes up to a multiple of 64
60 +# (XPUPlatform.update_block_size_for_backend), so use one.
61 +BLOCK_SIZE = 576
62 DEVICE_TYPE = current_platform.device_type
63 NUM_HIDDEN_LAYERS = 1
64 cur_step_action_idx = 0
@@ -541,26 +544,27 @@ def get_mamba_prefix_cache_step_configs(
544 async_scheduling: bool = False,
545 ) -> dict[str, TestConfig]:
546 a = async_scheduling
547 + B = BLOCK_SIZE
548 tests = {
549 "accept_1": TestConfig(
546 - num_prompt_tokens=554,
550 + num_prompt_tokens=B - 6,
551 num_generated_tokens=20,
552 num_accepted_tokens=1,
553 step_actions=[
550 - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
551 - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
554 + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
555 + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
556 StepAction(
553 - 555, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
557 + B - 5, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
558 ),
559 StepAction(
556 - 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
560 + B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
561 ),
558 - StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
559 - StepAction(558, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
560 - StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
561 - StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
562 + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
563 + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
564 + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
565 + StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
566 StepAction(
563 - 561,
567 + B + 1,
568 4,
569 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
570 (-1, -1),
@@ -570,19 +574,19 @@ def get_mamba_prefix_cache_step_configs(
574 ),
575 # test case 2.1: no hit, accept 2 tokens
576 "accept_2_1": TestConfig(
573 - num_prompt_tokens=554,
577 + num_prompt_tokens=B - 6,
578 num_generated_tokens=20,
579 num_accepted_tokens=2,
580 step_actions=[
577 - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
578 - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
581 + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
582 + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
583 StepAction(
580 - 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
584 + B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
585 ),
582 - StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
583 - StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
586 + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
587 + StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
588 StepAction(
585 - 562,
589 + B + 2,
590 4,
591 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
592 (-1, -1),
@@ -592,56 +596,56 @@ def get_mamba_prefix_cache_step_configs(
596 ),
597 # test case 2.2: no hit, accept 2 tokens
598 "accept_2_2": TestConfig(
595 - num_prompt_tokens=555,
599 + num_prompt_tokens=B - 5,
600 num_generated_tokens=20,
601 num_accepted_tokens=2,
602 step_actions=[
599 - StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
600 - StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
601 - StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
602 - StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
603 + StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
604 + StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
605 + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
606 + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
607 StepAction(
604 - 561,
608 + B + 1,
609 4,
610 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
611 (-1, -1),
612 (-1, -1),
613 ),
610 - StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
614 + StepAction(B + 3, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
615 ],
616 ),
617 "accept_3_1": TestConfig(
614 - num_prompt_tokens=553,
618 + num_prompt_tokens=B - 7,
619 num_generated_tokens=20,
620 num_accepted_tokens=3,
621 step_actions=[
618 - StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
619 - StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
622 + StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)),
623 + StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
624 StepAction(
621 - 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
625 + B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
626 ),
623 - StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
627 + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
628 StepAction(
625 - 562,
629 + B + 2,
630 4,
631 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
632 (-1, -1),
633 (-1, -1),
634 ),
631 - StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
635 + StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
636 ],
637 ),
638 "accept_3_2": TestConfig(
635 - num_prompt_tokens=554,
639 + num_prompt_tokens=B - 6,
640 num_generated_tokens=20,
641 num_accepted_tokens=3,
642 step_actions=[
639 - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
640 - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
641 - StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
642 - StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
643 + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
644 + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
645 + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
646 + StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
647 StepAction(
644 - 563,
648 + B + 3,
649 4,
650 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
651 (-1, -1),
@@ -650,87 +654,87 @@ def get_mamba_prefix_cache_step_configs(
654 ],
655 ),
656 "accept_3_3": TestConfig(
653 - num_prompt_tokens=555,
657 + num_prompt_tokens=B - 5,
658 num_generated_tokens=20,
659 num_accepted_tokens=3,
660 step_actions=[
657 - StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
658 - StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
659 - StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
661 + StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
662 + StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
663 + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
664 StepAction(
661 - 561,
665 + B + 1,
666 4,
667 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
668 (-1, -1),
669 (-1, -1),
670 ),
667 - StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
671 + StepAction(B + 4, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
672 ],
673 ),
674 "accept_4_1": TestConfig(
671 - num_prompt_tokens=553,
675 + num_prompt_tokens=B - 7,
676 num_generated_tokens=20,
677 num_accepted_tokens=4,
678 step_actions=[
675 - StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
676 - StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
677 - StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
679 + StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)),
680 + StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
681 + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
682 StepAction(
679 - 561,
683 + B + 1,
684 4,
685 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
686 (-1, -1),
687 (-1, -1),
688 ),
685 - StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
689 + StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
690 ],
691 ),
692 "accept_4_2": TestConfig(
689 - num_prompt_tokens=554,
693 + num_prompt_tokens=B - 6,
694 num_generated_tokens=25,
695 num_accepted_tokens=4,
696 step_actions=[
693 - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
694 - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
695 - StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
697 + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
698 + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
699 + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
700 StepAction(
697 - 562,
701 + B + 2,
702 4,
703 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
704 (-1, -1),
705 (-1, -1),
706 ),
703 - StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
707 + StepAction(B + 6, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
708 ],
709 ),
710 "accept_4_3": TestConfig(
707 - num_prompt_tokens=555,
711 + num_prompt_tokens=B - 5,
712 num_generated_tokens=25,
713 num_accepted_tokens=4,
714 step_actions=[
711 - StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
712 - StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
713 - StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
715 + StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
716 + StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
717 + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
718 StepAction(
715 - 563,
719 + B + 3,
720 4,
721 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
722 (-1, -1),
723 (-1, -1),
724 ),
721 - StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
725 + StepAction(B + 7, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
726 ],
727 ),
728 "accept_4_4": TestConfig(
725 - num_prompt_tokens=556,
729 + num_prompt_tokens=B - 4,
730 num_generated_tokens=25,
731 num_accepted_tokens=4,
732 step_actions=[
729 - StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)),
730 - StepAction(556, 4, [1, 1, 1, 1], (-1, -1), (3, 0)),
731 - StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
733 + StepAction(0, B - 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
734 + StepAction(B - 4, 4, [1, 1, 1, 1], (-1, -1), (3, 0)),
735 + StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
736 StepAction(
733 - 564,
737 + B + 4,
738 4,
739 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
740 (-1, -1),
@@ -739,23 +743,23 @@ def get_mamba_prefix_cache_step_configs(
743 ],
744 ),
745 "prompt_block_size": TestConfig(
742 - num_prompt_tokens=560,
746 + num_prompt_tokens=B,
747 num_generated_tokens=10,
748 num_accepted_tokens=4,
749 step_actions=[
746 - StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
747 - StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
750 + StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
751 + StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
752 ],
753 ),
754 "prompt_2_block_size": TestConfig(
751 - num_prompt_tokens=560 * 2,
755 + num_prompt_tokens=B * 2,
756 num_generated_tokens=10,
757 num_accepted_tokens=4,
758 step_actions=[
755 - StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
756 - StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
759 + StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
760 + StepAction(B, B, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
761 StepAction(
758 - 560 * 2,
762 + B * 2,
763 4,
764 [1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1],
765 (1, 2),
@@ -764,14 +768,14 @@ def get_mamba_prefix_cache_step_configs(
768 ],
769 ),
770 "prompt_2_block_size_10": TestConfig(
767 - num_prompt_tokens=560 * 2 + 10,
771 + num_prompt_tokens=B * 2 + 10,
772 num_generated_tokens=10,
773 num_accepted_tokens=4,
774 step_actions=[
771 - StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
772 - StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
775 + StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
776 + StepAction(B, B + 10, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
777 StepAction(
774 - 560 * 2 + 10,
778 + B * 2 + 10,
779 4,
780 [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1],
781 (-1, -1),
@@ -780,14 +784,14 @@ def get_mamba_prefix_cache_step_configs(
784 ],
785 ),
786 "prompt_3_block_size": TestConfig(
783 - num_prompt_tokens=560 * 3,
787 + num_prompt_tokens=B * 3,
788 num_generated_tokens=10,
789 num_accepted_tokens=4,
790 step_actions=[
787 - StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
788 - StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
791 + StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
792 + StepAction(B * 2, B, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
793 StepAction(
790 - 560 * 3,
794 + B * 3,
795 4,
796 [0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1],
797 (2, 3),
@@ -796,14 +800,14 @@ def get_mamba_prefix_cache_step_configs(
800 ],
801 ),
802 "prompt_3_block_size_10": TestConfig(
799 - num_prompt_tokens=560 * 3 + 10,
803 + num_prompt_tokens=B * 3 + 10,
804 num_generated_tokens=10,
805 num_accepted_tokens=4,
806 step_actions=[
803 - StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
804 - StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
807 + StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
808 + StepAction(B * 2, B + 10, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
809 StepAction(
806 - 560 * 3 + 10,
810 + B * 3 + 10,
811 4,
812 [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1],
813 (-1, -1),
@@ -812,21 +816,21 @@ def get_mamba_prefix_cache_step_configs(
816 ],
817 ),
818 "prompt_10_block_size": TestConfig(
815 - num_prompt_tokens=560 * 10,
819 + num_prompt_tokens=B * 10,
820 num_generated_tokens=10,
821 num_accepted_tokens=4,
822 step_actions=[
819 - StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
823 + StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
824 StepAction(
821 - 560 * 5,
822 - 560 * 4,
825 + B * 5,
826 + B * 4,
827 [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
828 (4, 8),
829 (-1, -1),
830 ),
831 StepAction(
828 - 560 * 9,
829 - 560,
832 + B * 9,
833 + B,
834 [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1]
835 if a
836 else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
@@ -834,7 +838,7 @@ def get_mamba_prefix_cache_step_configs(
838 (-1, -1),
839 ),
840 StepAction(
837 - 560 * 10,
841 + B * 10,
842 4,
843 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
844 if a
@@ -845,21 +849,21 @@ def get_mamba_prefix_cache_step_configs(
849 ],
850 ),
851 "prompt_10_block_size_10": TestConfig(
848 - num_prompt_tokens=560 * 10 + 10,
852 + num_prompt_tokens=B * 10 + 10,
853 num_generated_tokens=10,
854 num_accepted_tokens=4,
855 step_actions=[
852 - StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
856 + StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
857 StepAction(
854 - 560 * 5,
855 - 560 * 4,
858 + B * 5,
859 + B * 4,
860 [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
861 (4, 8),
862 (-1, -1),
863 ),
864 StepAction(
861 - 560 * 9,
862 - 560 + 10,
865 + B * 9,
866 + B + 10,
867 [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1]
868 if a
869 else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
vllm/platforms/xpu.py
+7 -4
@@ -464,13 +464,16 @@ class XPUPlatform(Platform):
464 )
465 cache_config.block_size = new_block_size
466 logger.info(
467 - "[XPU]Setting attention block size to %d tokens to ensure multiple of %d, "
468 - "set mamba_page_size_padded to %d bytes accordingly, before was %d bytes.",
467 + "[XPU]Setting attention block size to %d tokens to ensure multiple of %d.",
468 new_block_size,
469 kernel_block_size,
471 - cache_config.mamba_page_size_padded,
472 - original_mamba_page_size_padded,
470 )
471 + if original_mamba_page_size_padded is not None:
472 + logger.info(
473 + "[XPU]Scaled mamba_page_size_padded from %d to %d bytes accordingly.",
474 + original_mamba_page_size_padded,
475 + cache_config.mamba_page_size_padded,
476 + )
477
478 @classmethod
479 def support_hybrid_kv_cache(cls) -> bool: