7cab4368f2 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -8,7 +8,7 @@ Kimi-K2.5 extends Kimi-K2 with vision support.

8

9 from collections.abc import Iterable, Mapping, Sequence

10 from dataclasses import dataclass

11 -from typing import Annotated, Any, Literal

11 +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal

12

13 import torch

14 from torch import nn

@@ -25,10 +25,19 @@ from vllm.model_executor.layers.quantization.compressed_tensors import (

25 from vllm.model_executor.models.interfaces import (

26 SupportsEagle,

27 SupportsEagle3,

28 + SupportsEncoderCudaGraph,

29 SupportsMultiModal,

30 SupportsPP,

31 SupportsQuant,

32 )

33 +

34 +if TYPE_CHECKING:

35 + from vllm.v1.worker.encoder_cudagraph_defs import (

36 + EncoderCudaGraphCaptureInputs,

37 + EncoderCudaGraphConfig,

38 + EncoderCudaGraphReplayBuffers,

39 + EncoderItemSpec,

40 + )

41 from vllm.model_executor.models.kimi_k25_vit import (

42 KimiK25MultiModalProjector,

43 MoonViT3dPretrainedModel,

@@ -294,6 +303,7 @@ class KimiK25ForConditionalGeneration(

303 SupportsQuant,

304 SupportsEagle,

305 SupportsEagle3,

306 + SupportsEncoderCudaGraph,

307 ):

308 """Kimi-K2.5 model for conditional generation.

309

@@ -303,6 +313,7 @@ class KimiK25ForConditionalGeneration(

313 """

314

315 supports_encoder_tp_data = True

316 + supports_encoder_cudagraph: ClassVar[Literal[True]] = True

317

318 hf_to_vllm_mapper = WeightsMapper(

319 orig_to_new_prefix={

@@ -470,6 +481,241 @@ class KimiK25ForConditionalGeneration(

481 def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:

482 return self.language_model.get_eagle3_aux_hidden_state_layers()

483

484 + # ------------------------------------------------------------------ #

485 + # SupportsEncoderCudaGraph protocol #

486 + # Image-only (t == 1). Video chunks (t > 1) fall back to eager. #

487 + # ------------------------------------------------------------------ #

488 +

489 + @property

490 + def _encoder_cudagraph_pad_totals(self) -> dict[int, int]:

491 + """Row count of each captured buffer set, keyed by cu_seqlens ptr."""

492 + totals = self.__dict__.get("_encoder_cg_pad_totals")

493 + if totals is None:

494 + totals = {}

495 + self.__dict__["_encoder_cg_pad_totals"] = totals

496 + return totals

497 +

498 + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig":

499 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig

500 +

501 + pad_totals = self._encoder_cudagraph_pad_totals

502 +

503 + def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None:

504 + # Varlen attention requires cu_seqlens[-1] to equal the number of

505 + # rows actually passed in. The captured buffers are sized for the

506 + # full token budget, so a smaller real batch has to be completed

507 + # with one trailing padding sequence; declaring fewer rows than the

508 + # buffer holds is undefined behaviour and returns NaN on FlashAttn.

509 + total = pad_totals.get(dst.data_ptr())

510 + n = min(src.shape[0], dst.shape[0])

511 + dst[:n].copy_(src[:n])

512 + dst[n:] = total if total is not None else src[-1]

513 +

514 + return EncoderCudaGraphConfig(

515 + modalities=["vision_chunk"],

516 + buffer_keys=[

517 + "pixel_values",

518 + "pos_embeds",

519 + "rope_freqs_cis",

520 + "cu_seqlens",

521 + "max_seqlen",

522 + "merge_gather_idx",

523 + ],

524 + out_hidden_size=self.config.text_config.hidden_size,

525 + padding_logics={"cu_seqlens": pad_cu_seqlens},

526 + )

527 +

528 + def get_encoder_cudagraph_budget_range(

529 + self,

530 + vllm_config: VllmConfig,

531 + ) -> tuple[int, int]:

532 + # Min: 64 output tokens (e.g. ~128×128 image with patch 14 + merge 2×2).

533 + min_budget = 64

534 + max_budget = min(

535 + vllm_config.scheduler_config.max_num_batched_tokens,

536 + vllm_config.model_config.max_model_len,

537 + )

538 + return (min_budget, max_budget)

539 +

540 + def _get_grid_thw_list(

541 + self,

542 + mm_kwargs: dict[str, Any],

543 + ) -> list[list[int]]:

544 + grid_thws = mm_kwargs["grid_thws"]

545 + if isinstance(grid_thws, torch.Tensor):

546 + return [[int(x) for x in row] for row in grid_thws.tolist()]

547 + return [[int(x) for x in row] for row in grid_thws]

548 +

549 + def get_encoder_cudagraph_item_specs(

550 + self,

551 + mm_kwargs: dict[str, Any],

552 + ) -> list["EncoderItemSpec"]:

553 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

554 +

555 + kh, kw = self.vision_tower.merge_kernel_size

556 + specs = []

557 + for t, h, w in self._get_grid_thw_list(mm_kwargs):

558 + if t != 1:

559 + # Video chunks not supported in encoder CUDA graph;

560 + # sentinel forces eager fallback via the manager.

561 + specs.append(EncoderItemSpec(input_size=t * h * w, output_tokens=2**30))

562 + else:

563 + specs.append(

564 + EncoderItemSpec(

565 + input_size=h * w,

566 + output_tokens=(h // kh) * (w // kw),

567 + )

568 + )

569 + return specs

570 +

571 + def select_encoder_cudagraph_items(

572 + self,

573 + mm_kwargs: dict[str, Any],

574 + indices: list[int],

575 + ) -> dict[str, Any]:

576 + grid_thw_list = self._get_grid_thw_list(mm_kwargs)

577 + pixel_values = mm_kwargs["pixel_values"]

578 +

579 + if len(indices) == 0:

580 + return {

581 + "pixel_values": pixel_values[:0],

582 + "grid_thws": pixel_values.new_zeros((0, 3), dtype=torch.long),

583 + }

584 +

585 + patch_counts = [t * h * w for t, h, w in grid_thw_list]

586 + cum = [0]

587 + for pc in patch_counts:

588 + cum.append(cum[-1] + pc)

589 +

590 + selected_pv = torch.cat(

591 + [pixel_values[cum[i] : cum[i + 1]] for i in indices], dim=0

592 + )

593 + selected_grid = torch.tensor(

594 + [grid_thw_list[i] for i in indices],

595 + dtype=torch.long,

596 + device=pixel_values.device,

597 + )

598 + return {"pixel_values": selected_pv, "grid_thws": selected_grid}

599 +

600 + def prepare_encoder_cudagraph_capture_inputs(

601 + self,

602 + token_budget: int,

603 + max_batch_size: int,

604 + max_frames_per_batch: int,

605 + device: torch.device,

606 + dtype: torch.dtype,

607 + path: str = "default",

608 + ) -> "EncoderCudaGraphCaptureInputs":

609 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphCaptureInputs

610 +

611 + kh, kw = self.vision_tower.merge_kernel_size

612 + # Output tokens per item in the dummy grid (ceiling so total >= budget).

613 + per_item_out = (token_budget + max_batch_size - 1) // max_batch_size

614 +

615 + # Fit within RoPE precomputed max dimensions.

616 + rope = self.vision_tower.encoder.rope_2d

617 + max_wo = rope.max_width // kw

618 + wo = min(per_item_out, max_wo)

619 + ho = (per_item_out + wo - 1) // wo

620 + assert ho * kh <= rope.max_height, (

621 + f"per_item_out={per_item_out} exceeds RoPE grid capacity "

622 + f"(max {(rope.max_height // kh) * (rope.max_width // kw)} tokens)"

623 + )

624 +

625 + grid_thw_list = [[1, ho * kh, wo * kw] for _ in range(max_batch_size)]

626 +

627 + ps = self.vision_tower.patch_size

628 + if isinstance(ps, int):

629 + ps = (ps, ps)

630 + total_patches = max_batch_size * ho * kh * wo * kw

631 + dummy_pixel_values = torch.zeros(

632 + total_patches, 3, ps[0], ps[1], device=device, dtype=dtype

633 + )

634 +

635 + # max_seqlen must cover the worst case: one item consuming the full

636 + # budget, i.e. token_budget * kh * kw patches.

637 + # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can append

638 + # a padding sequence covering rows the real batch does not fill.

639 + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata(

640 + grid_thw_list,

641 + max_batch_size=max_batch_size + 1,

642 + max_seqlen_override=token_budget * kh * kw,

643 + device=device,

644 + )

645 +

646 + values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values}

647 + values.update({k: v for k, v in metadata.items() if v is not None})

648 +

649 + cu_seqlens = values.get("cu_seqlens")

650 + if cu_seqlens is not None:

651 + self._encoder_cudagraph_pad_totals[cu_seqlens.data_ptr()] = total_patches

652 +

653 + return EncoderCudaGraphCaptureInputs(values=values)

654 +

655 + def prepare_encoder_cudagraph_replay_buffers(

656 + self,

657 + mm_kwargs: dict[str, Any],

658 + max_batch_size: int,

659 + max_frames_per_batch: int,

660 + path: str = "default",

661 + ) -> "EncoderCudaGraphReplayBuffers":

662 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

663 +

664 + grid_thw_list = self._get_grid_thw_list(mm_kwargs)

665 + pixel_values = mm_kwargs["pixel_values"]

666 +

667 + # Unpadded: pad_cu_seqlens completes the tail with the padding sequence

668 + # so cu_seqlens[-1] matches the captured buffer's row count.

669 + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata(

670 + grid_thw_list,

671 + max_batch_size=None,

672 + device=pixel_values.device,

673 + )

674 +

675 + values: dict[str, torch.Tensor | None] = {"pixel_values": pixel_values}

676 + values.update(metadata)

677 + return EncoderCudaGraphReplayBuffers(values=values)

678 +

679 + def encoder_cudagraph_forward(

680 + self,

681 + inputs: dict[str, torch.Tensor],

682 + path: str = "default",

683 + ) -> torch.Tensor:

684 + pixel_values = inputs.pop("pixel_values")

685 + # Remaining keys (pos_embeds, rope_freqs_cis, cu_seqlens, max_seqlen,

686 + # merge_gather_idx, sequence_lengths) are consumed as encoder_metadata.

687 + encoder_metadata = inputs

688 +

689 + # Fast path: uses precomputed pos_embeds + merge_gather_idx.

690 + vt_output = self.vision_tower(

691 + pixel_values, grid_thws=None, encoder_metadata=encoder_metadata

692 + )

693 + # vt_output: (total_output_tokens, kh*kw, vit_hidden_dim)

694 +

695 + proj_dtype = next(self.mm_projector.parameters()).dtype

696 + if vt_output.dtype != proj_dtype:

697 + vt_output = vt_output.to(proj_dtype)

698 +

699 + projected = self.mm_projector(vt_output)

700 + return projected.view(-1, self.config.text_config.hidden_size)

701 +

702 + def encoder_eager_forward(

703 + self,

704 + mm_kwargs: dict[str, Any],

705 + path: str = "default",

706 + ) -> torch.Tensor:

707 + media_input = self._parse_and_validate_media_input(**mm_kwargs)

708 + if media_input is None:

709 + proj_dtype = next(self.mm_projector.parameters()).dtype

710 + return torch.zeros(

711 + 0,

712 + self.config.text_config.hidden_size,

713 + device=self.device,

714 + dtype=proj_dtype,

715 + )

716 + embeddings = self._process_media_input(media_input)

717 + return torch.cat(embeddings, dim=0)

718 +

719 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

720 loader = AutoWeightsLoader(self)

721 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

[MM][CG] Support ViT full CUDA graph for Kimi-K2.5 (#50929)