@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a0e0efd6bd

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

[Model] Support DP for ViT on Kimi-VL-A3B-Thinking-2506 (#23817)

Signed-off-by: Junhong <liujunhong11@huawei.com> Signed-off-by: LJH-LBJ <98734602+LJH-LBJ@users.noreply.github.com> Co-authored-by: Junhong <liujunhong11@huawei.com> Co-authored-by: LJH-LBJ <98734602+LJH-LBJ@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com>

WeiQing Chen committed Sep 2, 2025 at 00:56 UTC a0e0efd6bdcfd071e5f4123319628887dfd4973d

6 files changed +157 -62

docs/configuration/optimization.md

@@ -174,6 +174,7 @@ Regardless, you need to set `mm_encoder_tp_mode="data"` in engine arguments to u

174

175 Known supported models:

176

177 +- Kimi-VL (<gh-pr:23817>)

178 - Llama4 (<gh-pr:18368>)

179 - MiniCPM-V-2.5 or above (<gh-pr:23327>, <gh-pr:23948>)

180 - Qwen2.5-VL (<gh-pr:22742>)

tests/multimodal/test_utils.py

+12 -6

@@ -636,8 +636,10 @@ def run_dp_sharded_mrope_vision_model_vs_direct(local_rank: int,

636

637 # Run the model through the sharded function

638 with torch.inference_mode():

639 - sharded_output = run_dp_sharded_mrope_vision_model(

640 - vision_model, pixel_values, grid_thw_list)

639 + sharded_output = run_dp_sharded_mrope_vision_model(vision_model,

640 + pixel_values,

641 + grid_thw_list,

642 + rope_type="rope_3d")

643 sharded_output = torch.cat(sharded_output, dim=0)

644

645 # Check that the world size is setup correctly

@@ -691,8 +693,10 @@ def run_dp_sharded_mrope_vision_model_empty_input_worker(

693

694 # Should handle empty input gracefully

695 with torch.inference_mode():

694 - output = run_dp_sharded_mrope_vision_model(vision_model, pixel_values,

695 - grid_thw_list)

696 + output = run_dp_sharded_mrope_vision_model(vision_model,

697 + pixel_values,

698 + grid_thw_list,

699 + rope_type="rope_3d")

700

701 assert len(output) == 0

702

@@ -745,8 +749,10 @@ def run_dp_sharded_mrope_vision_model_uneven_load_worker(

749

750 # Should handle uneven distribution without errors

751 with torch.inference_mode():

748 - output_tuple = run_dp_sharded_mrope_vision_model(

749 - vision_model, pixel_values, grid_thw_list)

752 + output_tuple = run_dp_sharded_mrope_vision_model(vision_model,

753 + pixel_values,

754 + grid_thw_list,

755 + rope_type="rope_3d")

756

757 # Verify output shape is reasonable

758 merge_factor = vision_model.spatial_merge_size**2

vllm/model_executor/models/kimi_vl.py

+39 -15

@@ -56,6 +56,7 @@ from transformers.activations import GELUActivation

56 from vllm.config import VllmConfig

57 from vllm.distributed import get_pp_group

58 from vllm.model_executor.layers.fused_moe import FusedMoE

59 +from vllm.model_executor.layers.linear import ReplicatedLinear

60 from vllm.model_executor.layers.logits_processor import LogitsProcessor

61 from vllm.model_executor.layers.vocab_parallel_embedding import (

62 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead)

@@ -76,6 +77,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,

77 BaseProcessingInfo, PromptReplacement,

78 PromptUpdate)

79 from vllm.multimodal.profiling import BaseDummyInputsBuilder

80 +from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model

81 from vllm.sequence import IntermediateTensors

82 from vllm.transformers_utils.configs import KimiVLConfig, MoonViTConfig

83 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config

@@ -93,8 +95,10 @@ class MaxImageTokenMeta:

96 class KimiVLMultiModalProjector(nn.Module):

96 - def __init__(self, config: KimiVLConfig):

98 + def __init__(self, config: KimiVLConfig, \

99 + use_data_parallel: bool = False, prefix: str = ""):

100 super().__init__()

101 + self.use_data_parallel = use_data_parallel

102

103 self.hidden_size = (config.vision_config.hidden_size *

104 config.vision_config.merge_kernel_size[0] *

@@ -102,20 +106,24 @@ class KimiVLMultiModalProjector(nn.Module):

106

107 self.pre_norm = torch.nn.LayerNorm(config.vision_config.hidden_size,

108 eps=1e-5)

105 - self.linear_1 = nn.Linear(self.hidden_size,

106 - self.hidden_size,

107 - bias=True)

109 + self.linear_1 = ReplicatedLinear(self.hidden_size,

110 + self.hidden_size,

111 + bias=True,

112 + prefix=maybe_prefix(

113 + prefix, "linear_1"))

114 + self.linear_2 = ReplicatedLinear(self.hidden_size,

115 + config.text_config.hidden_size,

116 + bias=True,

117 + prefix=maybe_prefix(

118 + prefix, "linear_2"))

119 self.act = GELUActivation()

109 - self.linear_2 = nn.Linear(self.hidden_size,

110 - config.text_config.hidden_size,

111 - bias=True)

120

121 def forward(self, image_features: torch.Tensor) -> torch.Tensor:

122 hidden_states = self.pre_norm(image_features).view(

123 -1, self.hidden_size)

116 - hidden_states = self.linear_1(hidden_states)

124 + hidden_states, _ = self.linear_1(hidden_states)

125 hidden_states = self.act(hidden_states)

118 - hidden_states = self.linear_2(hidden_states)

126 + hidden_states, _ = self.linear_2(hidden_states)

127 return hidden_states

128

129

@@ -273,6 +281,8 @@ class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):

281 class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

282 SupportsPP):

283

284 + supports_encoder_tp_data = True

285 +

286 @classmethod

287 def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

288 if modality.startswith("image"):

@@ -292,10 +302,17 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

302 quant_config = vllm_config.quant_config

303

304 assert isinstance(config.vision_config, MoonViTConfig)

295 -

296 - self.vision_tower = MoonVitPretrainedModel(config.vision_config)

297 -

298 - self.multi_modal_projector = KimiVLMultiModalProjector(config=config)

305 + self.use_data_parallel = model_config.multimodal_config.mm_encoder_tp_mode == "data"

306 + self.hidden_size = config.text_config.hidden_size

307 + self.vision_tower = MoonVitPretrainedModel(config.vision_config,

308 + self.use_data_parallel,

309 + prefix=maybe_prefix(

310 + prefix, "vision_tower"))

311 +

312 + self.multi_modal_projector = KimiVLMultiModalProjector(

313 + config=config,

314 + use_data_parallel=self.use_data_parallel,

315 + prefix=maybe_prefix(prefix, "multi_modal_projector"))

316

317 self.quant_config = quant_config

318 sub_vllm_config = copy.deepcopy(vllm_config)

@@ -376,13 +393,19 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

393

394 pixel_values = inputs["pixel_values"]

395 image_grid_hws = inputs["image_grid_hws"]

379 - return self.vision_tower(pixel_values, image_grid_hws)

396 + if self.use_data_parallel:

397 + return run_dp_sharded_mrope_vision_model(self.vision_tower,

398 + pixel_values,

399 + image_grid_hws.tolist(),

400 + rope_type="rope_2d")

401 + else:

402 + return self.vision_tower(pixel_values, image_grid_hws)

403

404 def _process_image_input(self,

405 image_input: KimiVLImageInputs) -> torch.Tensor:

406 assert image_input["type"] == "pixel_values"

407 image_features = self._process_image_pixels(image_input)

385 - assert isinstance(image_features, list)

408 + assert isinstance(image_features, (list, tuple))

409 lengths = [x.shape[0] for x in image_features]

410 return self.multi_modal_projector(

411 torch.cat(image_features)).split(lengths)

@@ -496,6 +519,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

519 expert_params_mapping = []

520

521 params_dict = dict(self.named_parameters())

522 +

523 for args in weights:

524 name, loaded_weight = args[:2]

525 kwargs = args[2] if len(args) > 2 else {}

vllm/model_executor/models/moonvit.py

+55 -22

@@ -42,7 +42,6 @@

42 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

43 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

44 # SOFTWARE.

45 -import math

45 from collections.abc import Sequence

46 from copy import deepcopy

47 from functools import cached_property

@@ -55,6 +54,8 @@ from transformers.activations import ACT2FN, PytorchGELUTanh

54 from transformers.modeling_utils import PreTrainedModel

55 from transformers.utils import is_flash_attn_2_available

57 +from vllm.model_executor.layers.linear import ReplicatedLinear

58 +from vllm.model_executor.models.utils import maybe_prefix

59 from vllm.transformers_utils.configs.moonvit import MoonViTConfig

61 if is_flash_attn_2_available():

@@ -383,21 +384,30 @@ class MLP2(nn.Module):

384 bias: whether to use bias in linear layer.

385 """

386

386 - def __init__(self, dims: list[int], activation, bias=True):

387 + def __init__(self,

388 + dims: list[int],

389 + activation,

390 + bias=True,

391 + prefix: str = "",

392 + use_data_parallel: bool = False):

393 super().__init__()

394 assert len(dims) == 3

389 - self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)

390 - self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)

395 + self.use_data_parallel = use_data_parallel

396 + self.fc0 = ReplicatedLinear(dims[0],

397 + dims[1],

398 + bias=bias,

399 + prefix=maybe_prefix(prefix, "fc0"))

400 + self.fc1 = ReplicatedLinear(dims[1],

401 + dims[2],

402 + bias=bias,

403 + prefix=maybe_prefix(prefix, "fc1"))

404 self.activation = activation

392 - for m in [self.fc0, self.fc1]:

393 - nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))

394 - if m.bias is not None:

395 - nn.init.zeros_(m.bias)

405

406 def forward(self, x: torch.Tensor) -> torch.Tensor:

398 - x = self.fc0(x)

407 + x, _ = self.fc0(x)

408 x = self.activation(x)

400 - return self.fc1(x)

409 + x, _ = self.fc1(x)

410 + return x

411

412

413 class MoonVitEncoderLayer(nn.Module):

@@ -407,6 +417,8 @@ class MoonVitEncoderLayer(nn.Module):

417 num_heads: int,

418 hidden_dim: int,

419 mlp_dim: int,

420 + prefix: str = "",

421 + use_data_parallel: bool = False,

422 *,

423 attn_implementation: str = "sdpa",

424 activation=F.gelu,

@@ -423,9 +435,19 @@ class MoonVitEncoderLayer(nn.Module):

435

436 self.norm0 = nn.LayerNorm(hidden_dim)

437 self.norm1 = nn.LayerNorm(hidden_dim)

426 - self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)

427 - self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)

428 - self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)

438 + self.use_data_parallel = use_data_parallel

439 + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim],

440 + activation,

441 + prefix=f"{prefix}.mlp",

442 + use_data_parallel=use_data_parallel)

443 + self.wqkv = ReplicatedLinear(hidden_dim,

444 + hidden_dim * 3,

445 + bias=attn_bias,

446 + prefix=f"{prefix}.wqkv")

447 + self.wo = ReplicatedLinear(hidden_dim,

448 + hidden_dim,

449 + bias=attn_bias,

450 + prefix=f"{prefix}.wo")

451

452 def attention_qkvpacked(

453 self,

@@ -438,7 +460,7 @@ class MoonVitEncoderLayer(nn.Module):

460 x (torch.Tensor): (batch_size, seqlen, hidden_dim)

461 cu_seqlens (torch.Tensor):

462 """

441 - xqkv = self.wqkv(x)

463 + xqkv, _ = self.wqkv(x)

464

465 qkv_shape = xqkv.size()[:-1] + (

466 3,

@@ -457,8 +479,7 @@ class MoonVitEncoderLayer(nn.Module):

479 xv,

480 q_cu_seqlens=cu_seqlens,

481 k_cu_seqlens=cu_seqlens)

460 -

461 - attn_out = self.wo(attn_out)

482 + attn_out, _ = self.wo(attn_out)

483 return attn_out

484

485 def forward(

@@ -494,13 +515,17 @@ class MoonVitEncoder(nn.Module):

515 hidden_dim: int,

516 num_layers: int,

517 block_cfg: dict,

518 + prefix: str = "",

519 + use_data_parallel: bool = False,

520 ) -> None:

521 super().__init__()

522

523 self.rope_2d = Rope2DPosEmb(

524 block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512)

525 self.blocks = nn.ModuleList(

503 - [MoonVitEncoderLayer(**block_cfg) for _ in range(num_layers)])

526 + [MoonVitEncoderLayer(use_data_parallel=use_data_parallel, \

527 + prefix=f"{prefix}.blocks.{layer_idx}", \

528 + **block_cfg) for layer_idx in range(num_layers)])

529 self.final_layernorm = nn.LayerNorm(hidden_dim)

530

531 def forward(self, hidden_states: torch.Tensor,

@@ -508,10 +533,9 @@ class MoonVitEncoder(nn.Module):

533 rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(

534 grid_hws=grid_hw)

535

511 - lengths = torch.cat((

512 - torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),

513 - grid_hw[:, 0] * grid_hw[:, 1],

514 - ))

536 + lengths = torch.cat(

537 + (torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),

538 + (grid_hw[:, 0] * grid_hw[:, 1]).to(hidden_states.device)))

539 cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)

540

541 for _, block in enumerate(self.blocks):

@@ -587,11 +611,19 @@ class MoonVitPretrainedModel(PreTrainedModel):

611 _supports_flash_attn_2 = True

612 _supports_sdpa = True

613

590 - def __init__(self, config: MoonViTConfig, *inputs, **kwargs):

614 + def __init__(self,

615 + config: MoonViTConfig,

616 + use_data_parallel: bool = False,

617 + prefix: str = "",

618 + *inputs,

619 + **kwargs):

620 super().__init__(config, *inputs, **kwargs)

621 config = deepcopy(config)

622 + self.use_data_parallel = use_data_parallel

623 self.merge_kernel_size = config.merge_kernel_size

624 + self.hidden_size = config.hidden_size

625 self.patch_size = config.patch_size

626 + self.vit_processing_type = "rope_2d"

627 self.patch_embed = MoonVisionPatchEmbed(

628 out_dim=config.hidden_size,

629 patch_size=config.patch_size,

@@ -610,6 +642,7 @@ class MoonVitPretrainedModel(PreTrainedModel):

642 "attn_bias": True,

643 "attn_implementation": config._attn_implementation,

644 },

645 + prefix=f"{prefix}.encoder",

646 )

647

648 def forward(self, pixel_values: torch.Tensor,

vllm/model_executor/models/qwen2_5_vl.py

+8 -4

@@ -1021,8 +1021,10 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1021 pixel_values = image_input["pixel_values"]

1022

1023 if self.use_data_parallel:

1024 - return run_dp_sharded_mrope_vision_model(

1025 - self.visual, pixel_values, grid_thw_list)

1024 + return run_dp_sharded_mrope_vision_model(self.visual,

1025 + pixel_values,

1026 + grid_thw_list,

1027 + rope_type="rope_3d")

1028 else:

1029 image_embeds = self.visual(pixel_values,

1030 grid_thw=grid_thw_list)

@@ -1048,8 +1050,10 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1050 else:

1051 pixel_values_videos = video_input["pixel_values_videos"]

1052 if self.use_data_parallel:

1051 - return run_dp_sharded_mrope_vision_model(

1052 - self.visual, pixel_values_videos, grid_thw_list)

1053 + return run_dp_sharded_mrope_vision_model(self.visual,

1054 + pixel_values_videos,

1055 + grid_thw_list,

1056 + rope_type="rope_3d")

1057 else:

1058 video_embeds = self.visual(pixel_values_videos,

1059 grid_thw=grid_thw_list)

vllm/multimodal/utils.py

+42 -15

@@ -9,7 +9,7 @@ from collections.abc import Iterable

9 from concurrent.futures import ThreadPoolExecutor

10 from itertools import groupby

11 from pathlib import Path

12 -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union

12 +from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar, Union

13 from urllib.parse import ParseResult, urlparse

14 from urllib.request import url2pathname

@@ -444,7 +444,6 @@ def run_dp_sharded_vision_model(image_input: torch.Tensor,

444 Args:

445 image_input (torch.Tensor): Image input tensor.

446 vision_model (torch.nn.Module): Vision model.

447 -

447 Returns:

448 torch.Tensor: Output image embeddings

449 """

@@ -542,6 +541,8 @@ def run_dp_sharded_mrope_vision_model(

541 vision_model: torch.nn.Module,

542 pixel_values: torch.Tensor,

543 grid_thw_list: list[list[int]],

544 + *,

545 + rope_type: Literal["rope_3d", "rope_2d"],

546 ) -> tuple[torch.Tensor, ...]:

547 """Run a vision model with data parallelism (DP) sharding.

548 The function will shard the input image tensor on the

@@ -552,6 +553,10 @@ def run_dp_sharded_mrope_vision_model(

553 vision_model (torch.nn.Module): Vision model.

554 pixel_values (torch.Tensor): Image/Video input tensor.

555 grid_thw_list: List of grid dimensions for each image

556 + rope_type: Type of rope used in the vision model.

557 + Different rope types have different dimension to do ViT.

558 + "rope_3d" for 3D rope (e.g., Qwen2.5-VL)

559 + "rope_2d" for 2D rope (e.g., Kimi-VL)

560 Returns:

561 torch.Tensor: Output image embeddings

562

@@ -605,8 +610,12 @@ def run_dp_sharded_mrope_vision_model(

610 device=pixel_values.device,

611 dtype=pixel_values.dtype)

612 # embed_dim_reduction_factor = 2 * 2

608 - embed_dim_reduction_factor = (vision_model.spatial_merge_size *

609 - vision_model.spatial_merge_size)

613 + if rope_type == "rope_2d":

614 + embed_dim_reduction_factor = (vision_model.merge_kernel_size[0] *

615 + vision_model.merge_kernel_size[1])

616 + else:

617 + embed_dim_reduction_factor = (vision_model.spatial_merge_size *

618 + vision_model.spatial_merge_size)

619

620 # Find the max length across all ranks

621 # The output embedding of every DP rank has to be

@@ -617,23 +626,42 @@ def run_dp_sharded_mrope_vision_model(

626 local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]

627

628 # Run the vision model on the local pixel_values_local

620 - if pixel_values_local.shape[0] > 0:

621 - image_embeds_local = vision_model(pixel_values_local,

622 - local_grid_thw_list)

629 + if rope_type == "rope_2d":

630 + if pixel_values_local.shape[0] > 0:

631 + image_embeds_local = vision_model(

632 + pixel_values_local, torch.tensor(local_grid_thw_list))

633 + if isinstance(image_embeds_local, list):

634 + image_embeds_local = torch.cat(image_embeds_local, dim=0)

635 + else:

636 + out_dim = getattr(vision_model.config, "hidden_size", None)

637 + image_embeds_local = torch.empty(

638 + (0, embed_dim_reduction_factor, out_dim),

639 + device=pixel_values.device,

640 + dtype=pixel_values.dtype)

641 else:

624 - # Handle empty case

625 - image_embeds_local = torch.empty((0, vision_model.out_hidden_size),

626 - device=pixel_values.device,

627 - dtype=pixel_values.dtype)

642 + if pixel_values_local.shape[0] > 0:

643 + image_embeds_local = vision_model(pixel_values_local,

644 + local_grid_thw_list)

645 + else:

646 + # Handle empty case

647 + image_embeds_local = torch.empty((0, vision_model.out_hidden_size),

648 + device=pixel_values.device,

649 + dtype=pixel_values.dtype)

650

651 # Pad the output based on max_len_per_rank

652 # for tensor_model_parallel_all_gather to work

653 current_len = image_embeds_local.shape[0]

654 if current_len < max_len_per_rank:

655 padding_size = max_len_per_rank - current_len

634 - padding = torch.empty((padding_size, image_embeds_local.shape[1]),

635 - dtype=image_embeds_local.dtype,

636 - device=image_embeds_local.device)

656 + if rope_type == "rope_2d":

657 + padding = torch.empty((padding_size, image_embeds_local.shape[1],

658 + image_embeds_local.shape[2]),

659 + dtype=image_embeds_local.dtype,

660 + device=image_embeds_local.device)

661 + else:

662 + padding = torch.empty((padding_size, image_embeds_local.shape[1]),

663 + dtype=image_embeds_local.dtype,

664 + device=image_embeds_local.device)

665 image_embeds_local_padded = torch.cat([image_embeds_local, padding],

666 dim=0)

667 else:

@@ -674,7 +702,6 @@ def run_dp_sharded_mrope_vision_model(

702 embed_start:embed_start + img_patches]

703 embed_start += img_patches

704 current_idx += count

677 -

705 out_embeddings = tuple(embed for embed in original_order_embeddings

706 if embed is not None)

707 assert len(out_embeddings) == len(