d6b3d39b6d · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -16,6 +16,7 @@ from flashinfer import (

16 from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache

17 from flashinfer.prefill import trtllm_batch_context_with_kv_cache

18 from flashinfer.utils import FP4Tensor

19 +from typing_extensions import override

20

21 from vllm import envs

22 from vllm.attention.backends.abstract import (

@@ -59,6 +60,7 @@ from vllm.v1.attention.backends.utils import (

60 split_decodes_and_prefills,

61 )

62 from vllm.v1.kv_cache_interface import AttentionSpec

63 +from vllm.v1.utils import CpuGpuBuffer

64

65 FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024

66

@@ -181,7 +183,6 @@ class BatchDCPPrefillWrapper:

183 paged_kv_indptr_cpu: torch.Tensor,

184 paged_kv_indices: torch.Tensor,

185 paged_kv_last_page_len_cpu: torch.Tensor,

184 - prefill_start: int,

186 page_size: int,

187 num_qo_heads: int,

188 dcp_world_size: int,

@@ -200,7 +201,7 @@ class BatchDCPPrefillWrapper:

201 qo_indptr_cpu,

202 paged_kv_indptr_cpu,

203 paged_kv_indices,

203 - paged_kv_last_page_len_cpu[prefill_start:],

204 + paged_kv_last_page_len_cpu,

205 num_qo_heads * dcp_world_size,

206 num_kv_heads,

207 head_dim,

@@ -380,40 +381,103 @@ class FlashInferBackend(AttentionBackend):

381

382

383 @dataclass

383 -class FlashInferMetadata:

384 - num_actual_tokens: int # Number of tokens excluding padding.

384 +class FIPrefill:

385 + """Metadata for the native FlashInfer prefill pathway (non-TRTLLM)."""

386

386 - # The data type of the query

387 - q_data_type: torch.dtype

387 + wrapper: BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper

388

389 - slot_mapping: torch.Tensor

389

391 - # For flashinfer trtllm batch decode

390 +@dataclass

391 +class FIDecode:

392 + """Metadata for the native FlashInfer decode pathway (non-TRTLLM)."""

393 +

394 + wrapper: BatchDecodeWithPagedKVCacheWrapper

395 +

396 +

397 +@dataclass

398 +class TRTLLMPrefill:

399 + """Metadata for the TRTLLM prefill pathway."""

400 +

401 + block_tables: torch.Tensor

402 + """

403 + The slice of the block table tensor corresponding *only* to prefill requests.

404 + Shape: [num_prefills, max_num_blocks_per_seq]

405 + """

406 +

407 + seq_lens: torch.Tensor

408 + """

409 + The slice of the sequence lengths tensor corresponding *only* to prefill requests.

410 + Shape: [num_prefills]

411 + """

412 +

413 + cum_seq_lens_q: torch.Tensor

414 + cum_seq_lens_kv: torch.Tensor

415 +

416 max_q_len: int

393 - max_q_len_prefill: int

417 + """

418 + The maximum query length *among prefill requests*.

419 + """

420 +

421 max_seq_len: int

422 + """The maximum sequence length for KV Cache."""

423 +

424 +

425 +@dataclass

426 +class TRTLLMDecode:

427 + """Metadata for the TRTLLM decode pathway."""

428 +

429 + block_tables: torch.Tensor

430 + """

431 + The slice of the block table tensor corresponding *only* to decode requests.

432 + Shape: [num_decodes, max_num_blocks_per_seq]

433 + """

434 +

435 seq_lens: torch.Tensor

396 - block_table_tensor: torch.Tensor

397 - prefill_use_trtllm: bool

398 - decode_use_trtllm: bool

436 + """

437 + The slice of the sequence lengths tensor corresponding *only* to decode requests.

438 + Shape: [num_decodes]

439 + """

440 +

441 + max_seq_len: int

442 + """The maximum sequence length for KV Cache."""

443 +

444 +

445 +@dataclass

446 +class FlashInferMetadata:

447 + num_actual_tokens: int

448 + """Total number of tokens in the batch (excluding padding)."""

449 +

450 + slot_mapping: torch.Tensor

451 + """Tensor for writing K/V to the cache. Shape: [num_actual_tokens]"""

452 +

453 + q_data_type: torch.dtype

454

400 - # For handling prefill decode split

455 num_decodes: int

456 num_decode_tokens: int

457 num_prefills: int

458 num_prefill_tokens: int

459

406 - # For cascade attention (CPU for planning).

407 - use_cascade: bool

460 + prefill: FIPrefill | TRTLLMPrefill | None

461 + """

462 + Holds the metadata for the prefill portion of the batch.

463 + Will be `None` if `num_prefill_tokens == 0`.

464 + """

465 +

466 + decode: FIDecode | TRTLLMDecode | None

467 + """

468 + Holds the metadata for the decode portion of the batch.

469 + Will be `None` if `num_decode_tokens == 0`.

470 + """

471

409 - prefill_wrapper: (

410 - BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None

411 - ) = None

412 - decode_wrapper: BatchDecodeWithPagedKVCacheWrapper | None = None

413 - cascade_wrapper: MultiLevelCascadeAttentionWrapper | None = None

472 + # --- Special Case: Cascade Attention ---

473

415 - qo_indptr_gpu: torch.Tensor | None = None

416 - paged_kv_indptr_gpu: torch.Tensor | None = None

474 + use_cascade: bool

475 + """

476 + If True, the entire batch is a cascade attention call, and the

477 + `prefill` and `decode` fields will both be None.

478 + """

479 +

480 + cascade_wrapper: MultiLevelCascadeAttentionWrapper | None

481

482

483 class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

@@ -482,6 +546,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

546 self.dcp_world_size = 1

547 self.dcp_rank = 0

548 self.dcp_kv_cache_interleave_size = 1

549 + self.use_dcp = self.dcp_world_size > 1

550

551 self.num_qo_heads = self.model_config.get_num_attention_heads(

552 self.vllm_config.parallel_config

@@ -535,34 +600,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

600 "sinks, please use trtllm on blackwell or flash attention on "

601 "earlier GPUs."

602 )

538 - # Preparing persistent buffers (device-side)

539 - self.paged_kv_indptr = torch.zeros(

540 - max_num_reqs + 1, dtype=torch.int32, device=self.device

541 - )

542 - self.paged_kv_indices = torch.zeros(

543 - max_num_pages, # max num pages possible

544 - dtype=torch.int32,

545 - device=self.device,

546 - )

547 - self.paged_kv_last_page_len = torch.zeros(

548 - max_num_reqs, dtype=torch.int32, device=self.device

549 - )

550 - # host-side buffer

551 - pin_memory = is_pin_memory_available()

552 - self.paged_kv_indptr_cpu = torch.zeros(

553 - max_num_reqs + 1, dtype=torch.int32, device="cpu", pin_memory=pin_memory

554 - )

555 - self.paged_kv_indptr_np = self.paged_kv_indptr_cpu.numpy()

556 - self.paged_kv_indptr_buffer = torch.zeros_like(

557 - self.paged_kv_indptr_cpu, pin_memory=pin_memory

558 - )

559 - self.paged_kv_indices_cpu = torch.zeros(

560 - max_num_pages, dtype=torch.int32, device="cpu", pin_memory=pin_memory

561 - )

562 - self.paged_kv_last_page_len_cpu = torch.zeros(

563 - max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=pin_memory

564 - )

565 - self.paged_kv_last_page_len_np = self.paged_kv_last_page_len_cpu.numpy()

603 + # Preparing persistent buffers

604 + self.pin_memory = is_pin_memory_available()

605 + self.paged_kv_indptr = self._make_buffer(max_num_reqs + 1)

606 + self.paged_kv_indptr_cpu_buffer = torch.zeros_like(

607 + self.paged_kv_indptr.cpu, pin_memory=self.pin_memory

608 + ) # Extra buffer for mutable paged_kv_indptr.cpu in cuda graph mode

609 + self.paged_kv_indices = self._make_buffer(max_num_pages)

610 + self.paged_kv_last_page_len = self._make_buffer(max_num_reqs)

611

612 if self.head_dim == 256 and current_platform.is_device_capability_family(100):

613 # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that

@@ -573,6 +618,18 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

618 "passing --block-size 32 or --block-size 64."

619 )

620

621 + def _make_buffer(

622 + self, *size: int | torch.SymInt, dtype: torch.dtype = torch.int32

623 + ) -> CpuGpuBuffer:

624 + return CpuGpuBuffer(

625 + *size,

626 + dtype=dtype,

627 + device=self.device,

628 + pin_memory=self.pin_memory,

629 + with_numpy=True,

630 + )

631 +

632 + @override # type: ignore[misc]

633 @classmethod

634 def get_cudagraph_support(

635 cls: type["FlashInferMetadataBuilder"],

@@ -607,7 +664,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

664 self,

665 ) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:

666 if self._prefill_wrapper is None:

610 - if self.dcp_world_size > 1:

667 + if self.use_dcp:

668 self._prefill_wrapper = BatchDCPPrefillWrapper(

669 workspace_buffer=self._get_workspace_buffer(),

670 )

@@ -626,9 +683,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

683

684 if decode_wrapper is None:

685 if use_cudagraph:

629 - paged_kv_indptr = self.paged_kv_indptr[: batch_size + 1]

630 - paged_kv_indices = self.paged_kv_indices

631 - paged_kv_last_page_len = self.paged_kv_last_page_len[:batch_size]

686 + paged_kv_indptr = self.paged_kv_indptr.gpu[: batch_size + 1]

687 + paged_kv_indices = self.paged_kv_indices.gpu

688 + paged_kv_last_page_len = self.paged_kv_last_page_len.gpu[:batch_size]

689 else:

690 paged_kv_indptr = None

691 paged_kv_indices = None

@@ -661,99 +718,43 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

718 )

719 return self._cascade_wrapper

720

664 - def build(

721 + def _compute_flashinfer_kv_metadata(

722 self,

666 - common_prefix_len: int,

667 - common_attn_metadata: CommonAttentionMetadata,

668 - fast_build: bool = False,

669 - ) -> FlashInferMetadata:

670 - num_reqs = common_attn_metadata.num_reqs

671 - num_actual_tokens = common_attn_metadata.num_actual_tokens

672 - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

673 - split_decodes_and_prefills(

674 - common_attn_metadata,

675 - decode_threshold=self.reorder_batch_threshold,

676 - require_uniform=True,

677 - )

678 - )

679 -

680 - page_size = self.page_size

681 - max_q_len = common_attn_metadata.max_query_len

682 - max_seq_len = common_attn_metadata.max_seq_len

683 - seq_lens = common_attn_metadata.seq_lens

684 - seq_lens_cpu = common_attn_metadata.seq_lens_cpu

685 - block_table_tensor = common_attn_metadata.block_table_tensor

686 - qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu

687 -

688 - if self.dcp_world_size > 1:

689 - if num_prefills > 0:

690 - qo_indptr_prefill_cpu = (

691 - qo_indptr_cpu[num_decodes:] - qo_indptr_cpu[num_decodes]

692 - )

693 - query_lens_prefill_cpu = (

694 - qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]

695 - )

696 - seq_lens_cpu[num_decodes:] = (

697 - seq_lens_cpu[num_decodes:] - query_lens_prefill_cpu

698 - )

699 -

700 - seq_lens_cpu = get_dcp_local_seq_lens(

701 - seq_lens_cpu,

702 - self.dcp_world_size,

703 - self.dcp_rank,

704 - self.dcp_kv_cache_interleave_size,

705 - )

706 -

707 - seq_lens_np = seq_lens_cpu.numpy()

708 - num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size

709 -

710 - use_cascade = common_prefix_len > 0

711 - if use_cascade:

712 - # Grab the blocks of the shared prefix from the first request.

713 - assert common_prefix_len % page_size == 0

714 - num_common_kv_blocks = common_prefix_len // page_size

715 -

716 - # Create CPU versions directly for cascade (no GPU versions needed)

717 - shared_qo_indptr_cpu = torch.tensor(

718 - [0, num_actual_tokens], dtype=torch.int32, device="cpu"

719 - )

720 - shared_kv_page_indptr_cpu = torch.tensor(

721 - [0, num_common_kv_blocks], dtype=torch.int32, device="cpu"

722 - )

723 - shared_kv_page_indices_cpu = block_table_tensor[0, :num_common_kv_blocks]

724 - shared_kv_last_page_len_cpu = torch.tensor(

725 - [page_size], dtype=torch.int32, device="cpu"

726 - )

723 + num_blocks_np: np.ndarray,

724 + seq_lens_np: np.ndarray,

725 + block_table_tensor: torch.Tensor,

726 + num_reqs: int,

727 + page_size: int,

728 + ) -> torch.Tensor:

729 + """

730 + Compute paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len for FlashInfer

731 + attention.

732

728 - # Remove the blocks of the shared prefix from all requests.

729 - block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]

730 - num_blocks_np -= num_common_kv_blocks

731 - else:

732 - shared_qo_indptr_cpu = None

733 - shared_kv_page_indptr_cpu = None

734 - shared_kv_page_indices_cpu = None

735 - shared_kv_last_page_len_cpu = None

733 + Results are stored in self.paged_kv_indptr,

734 + self.paged_kv_indices, self.paged_kv_last_page_len buffers.

735

736 + Returns paged_kv_indices, a GPU tensor with shape [num_actual_pages].

737 + """

738 # write self.paged_kv_indptr_cpu inplace (0-index is always 0)

739 np.cumsum(

740 num_blocks_np,

741 dtype=np.int32,

741 - out=self.paged_kv_indptr_np[1 : num_reqs + 1],

742 + out=self.paged_kv_indptr.np[1 : num_reqs + 1],

743 )

744 # NOTE(woosuk): Because self.paged_kv_indptr_cpu can be modified

745 # after this line (e.g., for cuda graphs), we need to copy the data to

746 # self.paged_kv_indptr_buffer to avoid race condition.

746 - self.paged_kv_indptr_buffer[: num_reqs + 1] = self.paged_kv_indptr_cpu[

747 + self.paged_kv_indptr_cpu_buffer[: num_reqs + 1] = self.paged_kv_indptr.cpu[

748 : num_reqs + 1

749 ]

749 - paged_kv_indptr = self.paged_kv_indptr[: num_reqs + 1]

750 + paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1]

751 paged_kv_indptr.copy_(

751 - self.paged_kv_indptr_buffer[: num_reqs + 1], non_blocking=True

752 + self.paged_kv_indptr_cpu_buffer[: num_reqs + 1], non_blocking=True

753 )

754

755 # write self.paged_kv_indices inplace

755 - num_actual_pages = self.paged_kv_indptr_np[num_reqs]

756 - paged_kv_indices = self.paged_kv_indices[:num_actual_pages]

756 + num_actual_pages = self.paged_kv_indptr.np[num_reqs]

757 + paged_kv_indices = self.paged_kv_indices.gpu[:num_actual_pages]

758 _copy_page_indices_kernel[(num_reqs,)](

759 paged_kv_indices,

760 block_table_tensor,

@@ -764,12 +765,41 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

765

766 # write self.paged_kv_last_page_len_cpu inplace

767 paged_kv_last_page_len_np = seq_lens_np % page_size

767 - self.paged_kv_last_page_len_np[:num_reqs] = np.where(

768 + self.paged_kv_last_page_len.np[:num_reqs] = np.where(

769 (paged_kv_last_page_len_np == 0) & (seq_lens_np != 0),

770 page_size,

771 paged_kv_last_page_len_np,

772 )

773 + return paged_kv_indices

774

775 + def build(

776 + self,

777 + common_prefix_len: int,

778 + common_attn_metadata: CommonAttentionMetadata,

779 + fast_build: bool = False,

780 + ) -> FlashInferMetadata:

781 + num_reqs = common_attn_metadata.num_reqs

782 + num_actual_tokens = common_attn_metadata.num_actual_tokens

783 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

784 + split_decodes_and_prefills(

785 + common_attn_metadata,

786 + decode_threshold=self.reorder_batch_threshold,

787 + require_uniform=True,

788 + )

789 + )

790 +

791 + page_size = self.page_size

792 + max_seq_len = common_attn_metadata.max_seq_len

793 + seq_lens = common_attn_metadata.seq_lens

794 + block_table_tensor = common_attn_metadata.block_table_tensor

795 + qo_indptr = common_attn_metadata.query_start_loc

796 + qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu

797 +

798 + # Step 1: Decide which dispatch modes to use:

799 + # - Cascade attention (distinct mode)

800 + # - Prefill (FI native or TRTLLM)

801 + # - Decode (FI native or TRTLLM)

802 + use_cascade = common_prefix_len > 0

803 uses_spec_reorder = self.reorder_batch_threshold > 1

804 prefill_use_trtllm = use_trtllm_attention(

805 self.num_qo_heads,

@@ -788,7 +818,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

818 self.use_trtllm_decode_attention and self.dcp_world_size <= 1

819 )

820

791 - if not (prefill_use_trtllm and decode_use_trtllm):

821 + all_uses_trtllm = (num_prefills == 0 or prefill_use_trtllm) and (

822 + num_decodes == 0 or decode_use_trtllm

823 + )

824 + is_only_trtllm_decode = num_prefills == 0 and (

825 + num_decodes > 0 and decode_use_trtllm

826 + )

827 +

828 + if not all_uses_trtllm:

829 if self.has_sinks:

830 raise NotImplementedError(

831 "FlashInfer backend currently does not support attention "

@@ -813,28 +850,102 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

850 # fall back to model dtype.

851 self.q_data_type = self.model_config.dtype

852

853 + # Step 2: Initialize the output metadata

854 + # Leave prefill/decode/cascade_wrapper empty, to be populated

855 + # case by case depending on the batch contents and backend selection.

856 attn_metadata = FlashInferMetadata(

857 num_actual_tokens=num_actual_tokens,

818 - q_data_type=self.q_data_type,

858 slot_mapping=common_attn_metadata.slot_mapping,

820 - max_q_len=max_q_len,

821 - max_q_len_prefill=max_q_len,

822 - max_seq_len=max_seq_len,

823 - seq_lens=seq_lens,

824 - block_table_tensor=block_table_tensor,

825 - prefill_use_trtllm=prefill_use_trtllm,

826 - decode_use_trtllm=decode_use_trtllm,

859 + q_data_type=self.q_data_type,

860 num_decodes=num_decodes,

861 num_decode_tokens=num_decode_tokens,

862 num_prefills=num_prefills,

863 num_prefill_tokens=num_prefill_tokens,

864 use_cascade=use_cascade,

865 + prefill=None,

866 + decode=None,

867 + cascade_wrapper=None,

868 )

869

834 - paged_kv_indptr_cpu = self.paged_kv_indptr_cpu[: 1 + num_reqs]

835 - paged_kv_last_page_len_cpu = self.paged_kv_last_page_len_cpu[:num_reqs]

870 + # Guard access to seq_lens_cpu, which may not always be needed

871 + # and can be expensive to retrieve in async mode.

872 + needs_seq_lens_cpu = self.use_dcp or use_cascade or not is_only_trtllm_decode

873 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu if needs_seq_lens_cpu else None

874 + seq_lens_np = seq_lens_cpu.numpy() if seq_lens_cpu is not None else None

875 + num_blocks_np = (

876 + (seq_lens_np + (page_size - 1)) // page_size

877 + if seq_lens_np is not None

878 + else None

879 + )

880 +

881 + # Adjust seq_lens_cpu for DCP

882 + if self.use_dcp:

883 + assert seq_lens_cpu is not None

884 + if num_prefills > 0:

885 + qo_indptr_prefill_cpu = (

886 + qo_indptr_cpu[num_decodes:] - qo_indptr_cpu[num_decodes]

887 + )

888 + query_lens_prefill_cpu = (

889 + qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]

890 + )

891 + seq_lens_cpu[num_decodes:] = (

892 + seq_lens_cpu[num_decodes:] - query_lens_prefill_cpu

893 + )

894 +

895 + seq_lens_cpu = get_dcp_local_seq_lens(

896 + seq_lens_cpu,

897 + self.dcp_world_size,

898 + self.dcp_rank,

899 + self.dcp_kv_cache_interleave_size,

900 + )

901 +

902 + # Adjust num_block_np for cascade attention

903 + if use_cascade:

904 + assert num_blocks_np is not None

905 + assert common_prefix_len % page_size == 0

906 + num_common_kv_blocks = common_prefix_len // page_size

907 + num_blocks_np -= num_common_kv_blocks

908 +

909 + # Compute paged_kv_indices if necessary

910 + needs_paged_kv_indices = use_cascade or not is_only_trtllm_decode

911 + if needs_paged_kv_indices:

912 + assert num_blocks_np is not None

913 + assert seq_lens_np is not None

914 + paged_kv_indices = self._compute_flashinfer_kv_metadata(

915 + num_blocks_np,

916 + seq_lens_np,

917 + block_table_tensor,

918 + num_reqs,

919 + page_size,

920 + )

921 + else:

922 + paged_kv_indices = None

923 +

924 + # Early-out for cascade attention

925 + if use_cascade:

926 + # Grab the blocks of the shared prefix from the first request.

927 + num_common_kv_blocks = common_prefix_len // page_size

928 +

929 + # Create CPU versions directly for cascade (no GPU versions needed)

930 + shared_qo_indptr_cpu = torch.tensor(

931 + [0, num_actual_tokens], dtype=torch.int32, device="cpu"

932 + )

933 + shared_kv_page_indptr_cpu = torch.tensor(

934 + [0, num_common_kv_blocks], dtype=torch.int32, device="cpu"

935 + )

936 + shared_kv_page_indices_cpu = block_table_tensor[0, :num_common_kv_blocks]

937 + shared_kv_last_page_len_cpu = torch.tensor(

938 + [page_size], dtype=torch.int32, device="cpu"

939 + )

940 +

941 + # Remove the blocks of the shared prefix from all requests.

942 + block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]

943 + num_blocks_np -= num_common_kv_blocks

944 +

945 + assert paged_kv_indices is not None

946 + paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: 1 + num_reqs]

947 + paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]

948

837 - if attn_metadata.use_cascade:

949 attn_metadata.cascade_wrapper = self._get_cascade_wrapper()

950 attn_metadata.cascade_wrapper.plan(

951 [shared_qo_indptr_cpu, qo_indptr_cpu],

@@ -852,91 +963,107 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

963 q_data_type=self.q_data_type,

964 kv_data_type=self.kv_cache_dtype,

965 )

855 - else:

856 - # Regular attention (common case).

857 - # Decodes are at the front and prefills are at the back.

858 - num_prefills = attn_metadata.num_prefills

859 - num_decodes = attn_metadata.num_decodes

860 - if num_prefills > 0:

861 - # Decodes are first so prefills start after the last decode

862 - prefill_start = num_decodes

863 - attn_metadata.prefill_wrapper = self._get_prefill_wrapper()

864 - assert qo_indptr_cpu[prefill_start:].shape[0] == num_prefills + 1

865 - assert paged_kv_indptr_cpu[prefill_start:].shape[0] == num_prefills + 1

866 - assert (

867 - paged_kv_last_page_len_cpu[prefill_start:].shape[0] == num_prefills

966 + return attn_metadata

967 +

968 + # Step 3: Handle prefill and decode pathways case by case

969 + ## PREFILL PATHWAY

970 + if num_prefills > 0:

971 + # Slices for shared prefill metadata

972 + prefill_start = num_decodes

973 + qo_indptr_prefill_cpu = (

974 + qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[prefill_start]

975 + )

976 + assert qo_indptr_prefill_cpu.shape[0] == num_prefills + 1

977 +

978 + if prefill_use_trtllm:

979 + # Create GPU versions

980 + qo_indptr_prefill_gpu = (

981 + qo_indptr[prefill_start:] - qo_indptr[prefill_start]

982 + )

983 + paged_kv_indptr_prefill_gpu = self.paged_kv_indptr.gpu[

984 + prefill_start : num_reqs + 1

985 + ]

986 + # Compute max_q_len for prefill requests

987 + query_lens_prefill_cpu = (

988 + qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]

989 )

869 - # Since prefill_wrapper.run() will be called with

870 - # query[num_decode_tokens:] we need to adjust the qo_indptr

871 - # to be relative to the start of the prefill queries.

872 - qo_indptr_cpu = (

873 - qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[prefill_start]

990 + max_q_len_prefill = int(query_lens_prefill_cpu.max().item())

991 + attn_metadata.prefill = TRTLLMPrefill(

992 + block_tables=block_table_tensor[prefill_start:],

993 + seq_lens=seq_lens[prefill_start:],

994 + cum_seq_lens_q=qo_indptr_prefill_gpu,

995 + cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,

996 + max_q_len=max_q_len_prefill,

997 + max_seq_len=max_seq_len,

998 )

875 - paged_kv_indptr_cpu = paged_kv_indptr_cpu[prefill_start:]

876 -

877 - # Recompute max_q_len for the slice of requests we are using

878 - # for prefills. This can be different from max_q_len when

879 - # we have a non-uniform batch with some short decodes offloaded

880 - # to the prefill pathway

881 - query_lens_prefill = qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]

882 - attn_metadata.max_q_len_prefill = int(query_lens_prefill.max().item())

883 -

884 - if not attn_metadata.prefill_use_trtllm:

885 - if self.dcp_world_size > 1:

886 - assert isinstance(

887 - attn_metadata.prefill_wrapper, BatchDCPPrefillWrapper

888 - )

889 - attn_metadata.prefill_wrapper.plan(

890 - qo_indptr_cpu=qo_indptr_cpu,

891 - paged_kv_indptr_cpu=paged_kv_indptr_cpu,

892 - paged_kv_indices=paged_kv_indices,

893 - paged_kv_last_page_len_cpu=paged_kv_last_page_len_cpu,

894 - prefill_start=prefill_start,

895 - page_size=self.page_size,

896 - num_qo_heads=self.num_qo_heads,

897 - dcp_world_size=self.dcp_world_size,

898 - num_kv_heads=self.num_kv_heads,

899 - head_dim=self.head_dim,

900 - sm_scale=self.sm_scale,

901 - window_left=self.window_left,

902 - logits_soft_cap=self.logits_soft_cap,

903 - q_data_type=self.q_data_type,

904 - kv_cache_dtype=self.kv_cache_dtype,

905 - prefill_fixed_split_size=self.prefill_fixed_split_size,

906 - disable_split_kv=self.disable_split_kv,

907 - )

908 - else:

909 - assert isinstance(

910 - attn_metadata.prefill_wrapper,

911 - BatchPrefillWithPagedKVCacheWrapper,

912 - )

913 - attn_metadata.prefill_wrapper.plan(

914 - qo_indptr_cpu,

915 - paged_kv_indptr_cpu,

916 - paged_kv_indices,

917 - paged_kv_last_page_len_cpu[prefill_start:],

918 - self.num_qo_heads,

919 - self.num_kv_heads,

920 - self.head_dim,

921 - self.page_size,

922 - causal=True,

923 - sm_scale=self.sm_scale,

924 - window_left=self.window_left,

925 - logits_soft_cap=self.logits_soft_cap,

926 - q_data_type=self.q_data_type,

927 - kv_data_type=self.kv_cache_dtype,

928 - fixed_split_size=self.prefill_fixed_split_size,

929 - disable_split_kv=self.disable_split_kv,

930 - )

999 + else:

1000 + prefill_wrapper = self._get_prefill_wrapper()

1001 + # Slicing CPU buffers that are only needed for FI native prefills

1002 + paged_kv_last_page_len_prefill_cpu = self.paged_kv_last_page_len.cpu[

1003 + prefill_start:num_reqs

1004 + ]

1005 + assert paged_kv_last_page_len_prefill_cpu.shape[0] == num_prefills

1006 + paged_kv_indptr_prefill_cpu = self.paged_kv_indptr.cpu[

1007 + prefill_start : num_reqs + 1

1008 + ]

1009 + assert paged_kv_indptr_prefill_cpu.shape[0] == num_prefills + 1

1010 + if self.use_dcp:

1011 + assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)

1012 + prefill_wrapper.plan(

1013 + qo_indptr_cpu=qo_indptr_prefill_cpu,

1014 + paged_kv_indptr_cpu=paged_kv_indptr_prefill_cpu,

1015 + paged_kv_indices=paged_kv_indices,

1016 + paged_kv_last_page_len_cpu=paged_kv_last_page_len_prefill_cpu,

1017 + page_size=self.page_size,

1018 + num_qo_heads=self.num_qo_heads,

1019 + dcp_world_size=self.dcp_world_size,

1020 + num_kv_heads=self.num_kv_heads,

1021 + head_dim=self.head_dim,

1022 + sm_scale=self.sm_scale,

1023 + window_left=self.window_left,

1024 + logits_soft_cap=self.logits_soft_cap,

1025 + q_data_type=self.q_data_type,

1026 + kv_cache_dtype=self.kv_cache_dtype,

1027 + prefill_fixed_split_size=self.prefill_fixed_split_size,

1028 + disable_split_kv=self.disable_split_kv,

1029 + )

1030 else:

932 - attn_metadata.qo_indptr_gpu = qo_indptr_cpu.to(

933 - self.device, non_blocking=True

1031 + assert isinstance(

1032 + prefill_wrapper,

1033 + BatchPrefillWithPagedKVCacheWrapper,

1034 )

935 - attn_metadata.paged_kv_indptr_gpu = paged_kv_indptr_cpu.to(

936 - self.device, non_blocking=True

1035 + prefill_wrapper.plan(

1036 + qo_indptr_prefill_cpu,

1037 + paged_kv_indptr_prefill_cpu,

1038 + paged_kv_indices,

1039 + paged_kv_last_page_len_prefill_cpu,

1040 + self.num_qo_heads,

1041 + self.num_kv_heads,

1042 + self.head_dim,

1043 + self.page_size,

1044 + causal=True,

1045 + sm_scale=self.sm_scale,

1046 + window_left=self.window_left,

1047 + logits_soft_cap=self.logits_soft_cap,

1048 + q_data_type=self.q_data_type,

1049 + kv_data_type=self.kv_cache_dtype,

1050 + fixed_split_size=self.prefill_fixed_split_size,

1051 + disable_split_kv=self.disable_split_kv,

1052 )

1053 + attn_metadata.prefill = FIPrefill(wrapper=prefill_wrapper)

1054

939 - if num_decodes > 0:

1055 + ## DECODE PATHWAY

1056 + if num_decodes > 0:

1057 + if decode_use_trtllm:

1058 + assert num_decode_tokens % num_decodes == 0, (

1059 + "TRTLLM decode requires uniform query lengths per request."

1060 + )

1061 + attn_metadata.decode = TRTLLMDecode(

1062 + block_tables=block_table_tensor[:num_decodes],

1063 + seq_lens=seq_lens[:num_decodes],

1064 + max_seq_len=max_seq_len,

1065 + )

1066 + else:

1067 pure_decode = num_prefills == 0

1068 use_cudagraph = (

1069 self.enable_cuda_graph

@@ -945,33 +1072,33 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1072 )

1073 num_input_tokens = num_decode_tokens

1074

948 - attn_metadata.decode_wrapper = self._get_decode_wrapper(

1075 + decode_wrapper = self._get_decode_wrapper(

1076 num_input_tokens, use_cudagraph

1077 )

951 - if not attn_metadata.decode_use_trtllm:

952 - # Use the persistent buffer with padding length,

953 - # instead of the same address but chunked version

954 - # in atten_metadata when using cudagraph.

955 - fast_plan_decode(

956 - attn_metadata.decode_wrapper,

957 - self.paged_kv_indptr_cpu[: num_input_tokens + 1],

958 - paged_kv_indices,

959 - self.paged_kv_last_page_len_cpu[:num_input_tokens],

960 - seq_lens_cpu[:num_input_tokens],

961 - self.num_qo_heads * self.dcp_world_size,

962 - self.num_kv_heads,

963 - self.head_dim,

964 - self.page_size,

965 - # Disable flashinfer's pos encoding and use vllm's rope.

966 - pos_encoding_mode="NONE",

967 - sm_scale=self.sm_scale,

968 - window_left=self.window_left,

969 - logits_soft_cap=self.logits_soft_cap,

970 - q_data_type=self.q_data_type,

971 - kv_data_type=self.kv_cache_dtype,

972 - fixed_split_size=self.decode_fixed_split_size,

973 - disable_split_kv=self.disable_split_kv,

974 - )

1078 + # Use the persistent buffer with padding length,

1079 + # instead of the same address but chunked version

1080 + # in atten_metadata when using cudagraph.

1081 + fast_plan_decode(

1082 + decode_wrapper,

1083 + self.paged_kv_indptr.cpu[: num_input_tokens + 1],

1084 + paged_kv_indices,

1085 + self.paged_kv_last_page_len.cpu[:num_input_tokens],

1086 + seq_lens_cpu[:num_input_tokens],

1087 + self.num_qo_heads * self.dcp_world_size,

1088 + self.num_kv_heads,

1089 + self.head_dim,

1090 + self.page_size,

1091 + # Disable flashinfer's pos encoding and use vllm's rope.

1092 + pos_encoding_mode="NONE",

1093 + sm_scale=self.sm_scale,

1094 + window_left=self.window_left,

1095 + logits_soft_cap=self.logits_soft_cap,

1096 + q_data_type=self.q_data_type,

1097 + kv_data_type=self.kv_cache_dtype,

1098 + fixed_split_size=self.decode_fixed_split_size,

1099 + disable_split_kv=self.disable_split_kv,

1100 + )

1101 + attn_metadata.decode = FIDecode(wrapper=decode_wrapper)

1102 return attn_metadata

1103

1104 def use_cascade_attention(self, *args, **kwargs) -> bool:

@@ -1104,6 +1231,9 @@ class FlashInferImpl(AttentionImpl):

1231 if self.bmm2_scale is None:

1232 self.bmm2_scale = layer._v_scale_float

1233

1234 + prefill_use_trtllm = isinstance(attn_metadata.prefill, TRTLLMPrefill)

1235 + decode_use_trtllm = isinstance(attn_metadata.decode, TRTLLMDecode)

1236 +

1237 # The attn+quant fusion happens when output_scale is provided.

1238 if output_scale is None:

1239 assert output_block_scale is None, (

@@ -1113,8 +1243,8 @@ class FlashInferImpl(AttentionImpl):

1243 assert attn_metadata.q_data_type == FP8_DTYPE, (

1244 "Query must be FP8 when attn+quant fusion happened."

1245 )

1116 - assert (

1117 - attn_metadata.prefill_use_trtllm and attn_metadata.decode_use_trtllm

1246 + assert (attn_metadata.num_prefills == 0 or prefill_use_trtllm) and (

1247 + attn_metadata.num_decodes == 0 or decode_use_trtllm

1248 ), "Must use TRT-LLM attn"

1249

1250 if output.dtype == FP8_DTYPE:

@@ -1191,22 +1321,25 @@ class FlashInferImpl(AttentionImpl):

1321

1322 # When using spec decoding, num_decodes can be < num_decode_tokens

1323 # because some decode requests may have more than one query token.

1194 - num_decodes = attn_metadata.num_decodes

1324 num_decode_tokens = attn_metadata.num_decode_tokens

1325 num_prefill_tokens = attn_metadata.num_prefill_tokens

1326

1327 stride_order = FlashInferBackend.get_kv_cache_stride_order()

1328 kv_cache_permute = kv_cache.permute(*stride_order)

1329 +

1330 + use_dcp = self.dcp_world_size > 1

1331 +

1332 # Regular attention (common case).

1333 # Decodes are at the front and prefills are at the back.

1334 if num_prefill_tokens > 0:

1203 - prefill_wrapper = attn_metadata.prefill_wrapper

1335 prefill_query = query[num_decode_tokens:]

1336 assert prefill_query.shape[0] == num_prefill_tokens

1206 - assert prefill_wrapper is not None

1337

1208 - if not attn_metadata.prefill_use_trtllm:

1209 - if self.dcp_world_size > 1:

1338 + if not prefill_use_trtllm:

1339 + assert isinstance(attn_metadata.prefill, FIPrefill)

1340 + prefill_wrapper = attn_metadata.prefill.wrapper

1341 + assert prefill_wrapper is not None

1342 + if use_dcp:

1343 assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)

1344 assert prefill_wrapper._context._window_left == self.window_left

1345 assert prefill_wrapper._context._logits_soft_cap == (

@@ -1247,11 +1380,12 @@ class FlashInferImpl(AttentionImpl):

1380 out=output[num_decode_tokens:],

1381 )

1382 else:

1383 + assert isinstance(attn_metadata.prefill, TRTLLMPrefill)

1384 # prefill_query may be non-contiguous

1385 prefill_query = prefill_query.contiguous()

1386 workspace_buffer = _get_trtllm_gen_workspace_buffer()

1253 - block_tables_prefill = attn_metadata.block_table_tensor[num_decodes:]

1254 - seq_lens_prefill = attn_metadata.seq_lens[num_decodes:]

1387 + block_tables_prefill = attn_metadata.prefill.block_tables

1388 + seq_lens_prefill = attn_metadata.prefill.seq_lens

1389

1390 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND

1391 assert get_kv_cache_layout() == "HND"

@@ -1298,13 +1432,13 @@ class FlashInferImpl(AttentionImpl):

1432 workspace_buffer=workspace_buffer,

1433 block_tables=mock_block_table,

1434 seq_lens=seq_lens_prefill,

1301 - max_q_len=attn_metadata.max_q_len_prefill,

1302 - max_kv_len=attn_metadata.max_seq_len,

1435 + max_q_len=attn_metadata.prefill.max_q_len,

1436 + max_kv_len=attn_metadata.prefill.max_seq_len,

1437 bmm1_scale=self.bmm1_scale,

1438 bmm2_scale=self.bmm2_scale,

1439 batch_size=attn_metadata.num_prefills,

1306 - cum_seq_lens_q=attn_metadata.qo_indptr_gpu,

1307 - cum_seq_lens_kv=attn_metadata.paged_kv_indptr_gpu,

1440 + cum_seq_lens_q=attn_metadata.prefill.cum_seq_lens_q,

1441 + cum_seq_lens_kv=attn_metadata.prefill.cum_seq_lens_kv,

1442 window_left=self.window_left,

1443 sinks=self.sinks,

1444 o_sf_scale=self.o_sf_scale,

@@ -1312,17 +1446,18 @@ class FlashInferImpl(AttentionImpl):

1446 )

1447

1448 if num_decode_tokens > 0:

1315 - decode_wrapper = attn_metadata.decode_wrapper

1449 decode_query = query[:num_decode_tokens]

1450 assert decode_query.shape[0] == num_decode_tokens

1318 - assert decode_wrapper is not None

1451

1320 - if not attn_metadata.decode_use_trtllm:

1452 + if not decode_use_trtllm:

1453 + assert isinstance(attn_metadata.decode, FIDecode)

1454 + decode_wrapper = attn_metadata.decode.wrapper

1455 + assert decode_wrapper is not None

1456 assert decode_wrapper._window_left == self.window_left

1457 assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap or 0.0)

1458 assert decode_wrapper._sm_scale == self.scale

1459

1325 - if self.dcp_world_size > 1:

1460 + if use_dcp:

1461 decode_query = get_dcp_group().all_gather(

1462 decode_query.contiguous(), dim=-2

1463 )

@@ -1357,12 +1492,11 @@ class FlashInferImpl(AttentionImpl):

1492 )

1493 else:

1494 # decode_query may be non-contiguous

1495 + assert isinstance(attn_metadata.decode, TRTLLMDecode)

1496 decode_query = decode_query.contiguous()

1497 workspace_buffer = _get_trtllm_gen_workspace_buffer()

1362 - block_tables_decode = attn_metadata.block_table_tensor[

1363 - :num_decode_tokens

1364 - ]

1365 - seq_lens_decode = attn_metadata.seq_lens[:num_decode_tokens]

1498 + block_tables_decode = attn_metadata.decode.block_tables

1499 + seq_lens_decode = attn_metadata.decode.seq_lens

1500

1501 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND

1502 assert get_kv_cache_layout() == "HND"

@@ -1397,7 +1531,7 @@ class FlashInferImpl(AttentionImpl):

1531 workspace_buffer=workspace_buffer,

1532 block_tables=block_tables_decode,

1533 seq_lens=seq_lens_decode,

1400 - max_seq_len=attn_metadata.max_seq_len,

1534 + max_seq_len=attn_metadata.decode.max_seq_len,

1535 bmm1_scale=self.bmm1_scale,

1536 bmm2_scale=self.bmm2_scale,

1537 window_left=self.window_left,

[Cleanup] Refactor FlashInferMetadataBuilder (#29128)