70584f69b1 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,2354 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

7 +#

8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

9 +# and OPT implementations in this library. It has been modified from its

10 +# original forms to accommodate minor architectural differences compared

11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

12 +#

13 +# Licensed under the Apache License, Version 2.0 (the "License");

14 +# you may not use this file except in compliance with the License.

15 +# You may obtain a copy of the License at

16 +#

17 +# http://www.apache.org/licenses/LICENSE-2.0

18 +#

19 +# Unless required by applicable law or agreed to in writing, software

20 +# distributed under the License is distributed on an "AS IS" BASIS,

21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 +# See the License for the specific language governing permissions and

23 +# limitations under the License.

24 +"""Inference-only Cohere Compass model compatible with HuggingFace weights."""

25 +

26 +import math

27 +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence

28 +from functools import lru_cache, partial

29 +from itertools import islice

30 +from typing import Annotated, Any, Literal, TypeAlias

31 +

32 +import einops

33 +import numpy as np

34 +import torch

35 +import torch.nn as nn

36 +import torch.nn.functional as F

37 +from transformers import BatchFeature

38 +from transformers.models.cohere_compass import (

39 + CohereCompassConfig,

40 + CohereCompassProcessor,

41 +)

42 +from transformers.models.cohere_compass.configuration_cohere_compass import (

43 + CohereCompassVisionConfig,

44 +)

45 +from transformers.models.cohere_compass.image_processing_cohere_compass import (

46 + CohereCompassImageProcessor,

47 +)

48 +from transformers.models.cohere_compass.image_processing_cohere_compass import (

49 + smart_resize as image_smart_resize,

50 +)

51 +

52 +from vllm.compilation.decorators import support_torch_compile

53 +from vllm.config import CacheConfig, VllmConfig

54 +from vllm.config.multimodal import BaseDummyOptions

55 +from vllm.distributed import (

56 + get_pp_group,

57 + get_tensor_model_parallel_world_size,

58 + parallel_state,

59 +)

60 +from vllm.distributed import utils as dist_utils

61 +from vllm.inputs import ModalityData, MultiModalDataDict

62 +from vllm.logger import init_logger

63 +from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY

64 +from vllm.model_executor.layers.attention import Attention

65 +from vllm.model_executor.layers.attention.mm_encoder_attention import (

66 + MMEncoderAttention,

67 +)

68 +from vllm.model_executor.layers.conv import Conv3dLayer

69 +from vllm.model_executor.layers.linear import (

70 + ColumnParallelLinear,

71 + QKVParallelLinear,

72 + RowParallelLinear,

73 +)

74 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

75 +from vllm.model_executor.layers.quantization import QuantizationConfig

76 +from vllm.model_executor.layers.rotary_embedding import get_rope

77 +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb

78 +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

79 +from vllm.model_executor.models.module_mapping import MultiModelKeys

80 +from vllm.multimodal import MULTIMODAL_REGISTRY

81 +from vllm.multimodal.inputs import (

82 + ImageItem,

83 + MultiModalFeatureSpec,

84 + MultiModalFieldConfig,

85 + MultiModalKwargsItems,

86 +)

87 +from vllm.multimodal.parse import (

88 + DictEmbeddingItems,

89 + ImageSize,

90 + ModalityDataItems,

91 + MultiModalDataItems,

92 + MultiModalDataParser,

93 +)

94 +from vllm.multimodal.processing import (

95 + BaseDummyInputsBuilder,

96 + BaseMultiModalProcessor,

97 + BaseProcessingInfo,

98 + PromptReplacement,

99 + PromptUpdate,

100 +)

101 +from vllm.multimodal.video_prune.evs import (

102 + compute_mrope_for_media,

103 + recompute_mrope_positions,

104 +)

105 +from vllm.sequence import IntermediateTensors

106 +from vllm.tokenizers.registry import cached_tokenizer_from_config

107 +from vllm.triton_utils import HAS_TRITON, tl, triton

108 +from vllm.utils.math_utils import round_up

109 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

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

111 +

112 +from ...utils.torch_utils import async_tensor_h2d

113 +from .commandr import CohereForCausalLM, CohereMLP, LayerNorm

114 +from .interfaces import (

115 + MultiModalEmbeddings,

116 + SupportsEagle,

117 + SupportsEagle3,

118 + SupportsEncoderCudaGraph,

119 + SupportsLoRA,

120 + SupportsMRoPE,

121 + SupportsMultiModal,

122 + SupportsMultiModalPruning,

123 + SupportsPP,

124 + _require_is_multimodal,

125 +)

126 +from .utils import (

127 + AutoWeightsLoader,

128 + WeightsMapper,

129 + _merge_multimodal_embeddings,

130 + extract_layer_index,

131 + make_empty_intermediate_tensors_factory,

132 + make_layers,

133 + maybe_prefix,

134 +)

135 +from .vision import (

136 + get_fp8_padded_hidden_size,

137 + get_vit_attn_backend,

138 + is_vit_use_data_parallel,

139 + run_dp_sharded_mrope_vision_model,

140 +)

141 +

142 +logger = init_logger(__name__)

143 +

144 +# ---------------------------------------------------------------------------

145 +# Triton kernel: fused bilinear position-embedding interpolation

146 +# ---------------------------------------------------------------------------

147 +# Replaces many small eager-mode CUDA kernels with a single launch.

148 +# The spatial-merge reorder is baked into the index math so the output

149 +# is ready to be added to the patch embeddings directly.

150 +# ---------------------------------------------------------------------------

151 +

152 +if HAS_TRITON:

153 +

154 + @triton.jit

155 + def _bilinear_pos_embed_kernel(

156 + embed_ptr,

157 + output_ptr,

158 + H,

159 + W,

160 + h_scale,

161 + w_scale,

162 + NUM_GRID: tl.constexpr,

163 + M_SIZE: tl.constexpr,

164 + HIDDEN_DIM: tl.constexpr,

165 + BLOCK_D: tl.constexpr,

166 + ):

167 + """Fused bilinear pos-embed interpolation with spatial-merge reorder."""

168 + pid = tl.program_id(0)

169 + total_spatial = H * W

170 + spatial_idx = pid % total_spatial

171 +

172 + num_blocks_w = W // M_SIZE

173 + block_idx = spatial_idx // (M_SIZE * M_SIZE)

174 + local_idx = spatial_idx % (M_SIZE * M_SIZE)

175 + br = block_idx // num_blocks_w

176 + bc = block_idx % num_blocks_w

177 + lr = local_idx // M_SIZE

178 + lc = local_idx % M_SIZE

179 + row = br * M_SIZE + lr

180 + col = bc * M_SIZE + lc

181 +

182 + h_frac = row.to(tl.float32) * h_scale

183 + w_frac = col.to(tl.float32) * w_scale

184 +

185 + hf = tl.math.floor(h_frac).to(tl.int32)

186 + wf = tl.math.floor(w_frac).to(tl.int32)

187 + hc = tl.minimum(hf + 1, NUM_GRID - 1)

188 + wc = tl.minimum(wf + 1, NUM_GRID - 1)

189 +

190 + dh = h_frac - hf.to(tl.float32)

191 + dw = w_frac - wf.to(tl.float32)

192 + w11 = dh * dw

193 + w10 = dh - w11

194 + w01 = dw - w11

195 + w00 = 1.0 - dh - w01

196 +

197 + off00 = (hf * NUM_GRID + wf) * HIDDEN_DIM

198 + off01 = (hf * NUM_GRID + wc) * HIDDEN_DIM

199 + off10 = (hc * NUM_GRID + wf) * HIDDEN_DIM

200 + off11 = (hc * NUM_GRID + wc) * HIDDEN_DIM

201 + out_off = pid * HIDDEN_DIM

202 +

203 + # Cast weights to output dtype so the multiply-accumulate stays

204 + # in the same precision as the native PyTorch implementation.

205 + out_dtype = output_ptr.dtype.element_ty

206 + w00_c = w00.to(out_dtype)

207 + w01_c = w01.to(out_dtype)

208 + w10_c = w10.to(out_dtype)

209 + w11_c = w11.to(out_dtype)

210 +

211 + for d in tl.range(0, HIDDEN_DIM, BLOCK_D):

212 + cols = d + tl.arange(0, BLOCK_D)

213 + mask = cols < HIDDEN_DIM

214 +

215 + e00 = tl.load(embed_ptr + off00 + cols, mask=mask)

216 + e01 = tl.load(embed_ptr + off01 + cols, mask=mask)

217 + e10 = tl.load(embed_ptr + off10 + cols, mask=mask)

218 + e11 = tl.load(embed_ptr + off11 + cols, mask=mask)

219 +

220 + val = w00_c * e00 + w01_c * e01 + w10_c * e10 + w11_c * e11

221 +

222 + tl.store(output_ptr + out_off + cols, val, mask=mask)

223 +

224 + def triton_pos_embed_interpolate(

225 + embed_weight: torch.Tensor,

226 + t: int,

227 + h: int,

228 + w: int,

229 + num_grid_per_side: int,

230 + m_size: int,

231 + dtype: torch.dtype,

232 + ) -> torch.Tensor:

233 + """Launch the fused Triton kernel for one (t,h,w) grid.

234 +

235 + Returns a tensor of shape ``(t * h * w, hidden_dim)`` with the

236 + bilinearly-interpolated position embeddings in spatial-merge order.

237 + """

238 + assert h % m_size == 0 and w % m_size == 0, (

239 + f"h={h} and w={w} must be divisible by m_size={m_size}"

240 + )

241 + hidden_dim = embed_weight.shape[1]

242 + total_out = t * h * w

243 + output = torch.empty(

244 + total_out,

245 + hidden_dim,

246 + device=embed_weight.device,

247 + dtype=dtype,

248 + )

249 +

250 + h_scale = float(num_grid_per_side - 1) / float(h - 1) if h > 1 else 0.0

251 + w_scale = float(num_grid_per_side - 1) / float(w - 1) if w > 1 else 0.0

252 +

253 + BLOCK_D = triton.next_power_of_2(hidden_dim)

254 +

255 + _bilinear_pos_embed_kernel[(total_out,)](

256 + embed_weight,

257 + output,

258 + h,

259 + w,

260 + h_scale,

261 + w_scale,

262 + num_grid_per_side,

263 + m_size,

264 + hidden_dim,

265 + BLOCK_D,

266 + )

267 + return output

268 +

269 +

270 +def pos_embed_interpolate_native(

271 + embed_weight: torch.Tensor,

272 + t: int,

273 + h: int,

274 + w: int,

275 + num_grid_per_side: int,

276 + m_size: int,

277 + dtype: torch.dtype,

278 +) -> torch.Tensor:

279 + """Eager PyTorch bilinear position-embedding interpolation.

280 +

281 + Returns a tensor of shape ``(t * h * w, hidden_dim)`` with the

282 + bilinearly-interpolated position embeddings in spatial-merge order.

283 + """

284 + assert h % m_size == 0 and w % m_size == 0, (

285 + f"h={h} and w={w} must be divisible by m_size={m_size}"

286 + )

287 + hidden_dim = embed_weight.shape[1]

288 + device = embed_weight.device

289 +

290 + h_idxs = torch.linspace(

291 + 0,

292 + num_grid_per_side - 1,

293 + h,

294 + dtype=torch.float32,

295 + device=device,

296 + )

297 + w_idxs = torch.linspace(

298 + 0,

299 + num_grid_per_side - 1,

300 + w,

301 + dtype=torch.float32,

302 + device=device,

303 + )

304 +

305 + h_floor = h_idxs.to(torch.long)

306 + w_floor = w_idxs.to(torch.long)

307 + h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)

308 + w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)

309 +

310 + dh = h_idxs - h_floor

311 + dw = w_idxs - w_floor

312 +

313 + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")

314 + h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")

315 + h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")

316 +

317 + w11 = dh_grid * dw_grid

318 + w10 = dh_grid - w11

319 + w01 = dw_grid - w11

320 + w00 = 1 - dh_grid - w01

321 +

322 + h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid])

323 + w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid])

324 + h_grid_idx = h_grid * num_grid_per_side

325 +

326 + indices = (h_grid_idx + w_grid).reshape(4, -1)

327 + weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)

328 + weights = weights.to(dtype=dtype)

329 +

330 + embeds = embed_weight[indices]

331 + embeds *= weights

332 + combined = embeds.sum(dim=0)

333 +

334 + combined = combined.reshape(h // m_size, m_size, w // m_size, m_size, hidden_dim)

335 + combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim)

336 + repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim)

337 + return repeated.to(dtype=dtype)

338 +

339 +

340 +class Cohere_VLImagePixelInputs(TensorSchema):

341 + """

342 + Dimensions:

343 + - np: Number of patches

344 + - ni: Number of images

345 + - cps: Number of channels * patch_size * patch_size

346 +

347 + Historical context:

348 + - pixel_values shape: (num_patches, num_channels * patch_size *

349 + patch_size)

350 + - image_grid_thw shape: (num_images, 3) in (grid_t, grid_h, grid_w)

351 + format.

352 + """

353 +

354 + type: Literal["pixel_values"]

355 +

356 + pixel_values: Annotated[

357 + torch.Tensor,

358 + TensorShape("np", "cps"),

359 + ]

360 +

361 + image_grid_thw: Annotated[

362 + torch.Tensor,

363 + TensorShape("ni", 3),

364 + ]

365 +

366 +

367 +class Cohere_VLImageEmbeddingInputs(TensorSchema):

368 + """

369 + Dimensions:

370 + - nf: Number of image features

371 + - hs: Hidden size

372 + - ni: Number of images

373 +

374 + Historical context:

375 + - image_embeds shape: (num_image_features, hidden_size)

376 + - num_image_features varies based on the number and resolution of the

377 + images.

378 + - hidden_size must match the hidden size of language model backbone.

379 + - image_grid_thw shape: (num_images, 3) in (grid_t, grid_h, grid_w)

380 + format

381 + """

382 +

383 + type: Literal["image_embeds"]

384 +

385 + image_embeds: Annotated[

386 + torch.Tensor,

387 + TensorShape("nf", "hs"),

388 + ]

389 +

390 + image_grid_thw: Annotated[

391 + torch.Tensor,

392 + TensorShape("ni", 3),

393 + ]

394 +

395 +

396 +Cohere_VLImageInputs: TypeAlias = (

397 + Cohere_VLImagePixelInputs | Cohere_VLImageEmbeddingInputs

398 +)

399 +

400 +

401 +class Cohere_VisionAttention(nn.Module):

402 + def __init__(

403 + self,

404 + embed_dim: int,

405 + num_heads: int,

406 + projection_size: int,

407 + quant_config: QuantizationConfig | None = None,

408 + prefix: str = "",

409 + ) -> None:

410 + super().__init__()

411 + # Per attention head and per partition values.

412 + use_data_parallel = is_vit_use_data_parallel()

413 + self.tp_size = (

414 + 1

415 + if use_data_parallel

416 + else parallel_state.get_tensor_model_parallel_world_size()

417 + )

418 + self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

419 + self.hidden_size_per_attention_head = dist_utils.divide(

420 + projection_size, num_heads

421 + )

422 + self.num_attention_heads_per_partition = dist_utils.divide(

423 + num_heads, self.tp_size

424 + )

425 +

426 + self.qkv = QKVParallelLinear(

427 + hidden_size=embed_dim,

428 + head_size=self.hidden_size_per_attention_head,

429 + total_num_heads=num_heads,

430 + total_num_kv_heads=num_heads,

431 + bias=True,

432 + quant_config=quant_config,

433 + prefix=f"{prefix}.qkv",

434 + disable_tp=use_data_parallel,

435 + )

436 +

437 + self.proj = RowParallelLinear(

438 + input_size=projection_size,

439 + output_size=embed_dim,

440 + quant_config=quant_config,

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

442 + disable_tp=use_data_parallel,

443 + )

444 +

445 + self.attn = MMEncoderAttention(

446 + num_heads=self.num_attention_heads_per_partition,

447 + head_size=self.hidden_size_per_attention_head,

448 + scale=self.hidden_size_per_attention_head**-0.5,

449 + prefix=f"{prefix}.attn",

450 + )

451 +

452 + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)

453 +

454 + def forward(

455 + self,

456 + x: torch.Tensor,

457 + cu_seqlens: torch.Tensor,

458 + rotary_pos_emb_cos: torch.Tensor,

459 + rotary_pos_emb_sin: torch.Tensor,

460 + max_seqlen: torch.Tensor, # Only used for Flash Attention

461 + # Only used for FlashInfer CuDNN backend.

462 + sequence_lengths: torch.Tensor | None,

463 + ) -> torch.Tensor:

464 + # [s, b, c] --> [s, b, head * 3 * head_dim]

465 + x, _ = self.qkv(x)

466 + seq_len, batch_size, _ = x.shape

467 +

468 + qkv = einops.rearrange(

469 + x,

470 + "s b (three head head_dim) -> b s three head head_dim",

471 + three=3,

472 + head=self.num_attention_heads_per_partition,

473 + )

474 +

475 + if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:

476 + qk, v = qkv[:, :, :2], qkv[:, :, 2]

477 +

478 + qk_reshaped = einops.rearrange(

479 + qk, "b s two head head_dim -> (two b) s head head_dim", two=2

480 + )

481 + qk_reshaped = qk_reshaped.contiguous()

482 + qk_rotated = self.apply_rotary_emb(

483 + qk_reshaped,

484 + rotary_pos_emb_cos,

485 + rotary_pos_emb_sin,

486 + )

487 + qk_rotated = qk_rotated.view(

488 + 2,

489 + batch_size,

490 + seq_len,

491 + self.num_attention_heads_per_partition,

492 + self.hidden_size_per_attention_head,

493 + )

494 + q, k = qk_rotated.unbind(dim=0)

495 + else:

496 + q, k, v = qkv.unbind(dim=2)

497 +

498 + context_layer = self.attn(

499 + query=q,

500 + key=k,

501 + value=v,

502 + cu_seqlens=cu_seqlens,

503 + max_seqlen=max_seqlen,

504 + sequence_lengths=sequence_lengths,

505 + )

506 +

507 + context_layer = einops.rearrange(

508 + context_layer, "b s h d -> s b (h d)", b=batch_size

509 + ).contiguous()

510 +

511 + output, _ = self.proj(context_layer)

512 + return output

513 +

514 +

515 +class Cohere_VisionPatchEmbed(nn.Module):

516 + def __init__(

517 + self,

518 + patch_size: int = 14,

519 + temporal_patch_size: int = 2,

520 + in_channels: int = 3,

521 + hidden_size: int = 1152,

522 + ) -> None:

523 + super().__init__()

524 + self.patch_size = patch_size

525 + self.temporal_patch_size = temporal_patch_size

526 + self.hidden_size = hidden_size

527 +

528 + kernel_size = (temporal_patch_size, patch_size, patch_size)

529 + self.proj = Conv3dLayer(

530 + in_channels,

531 + hidden_size,

532 + kernel_size=kernel_size,

533 + stride=kernel_size,

534 + bias=True,

535 + )

536 +

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

538 + L, C = x.shape

539 + x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)

540 + x = self.proj(x).view(L, self.hidden_size)

541 + return x

542 +

543 +

544 +class Cohere_VisionMLP(nn.Module):

545 + def __init__(

546 + self,

547 + in_features: int,

548 + hidden_features: int,

549 + bias: bool = False,

550 + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

551 + quant_config: QuantizationConfig | None = None,

552 + prefix: str = "",

553 + ):

554 + super().__init__()

555 + use_data_parallel = is_vit_use_data_parallel()

556 + self.linear_fc1 = ColumnParallelLinear(

557 + in_features,

558 + hidden_features,

559 + bias=bias,

560 + quant_config=quant_config,

561 + return_bias=False,

562 + prefix=f"{prefix}.linear_fc1",

563 + disable_tp=use_data_parallel,

564 + )

565 + self.linear_fc2 = RowParallelLinear(

566 + hidden_features,

567 + in_features,

568 + bias=bias,

569 + quant_config=quant_config,

570 + return_bias=False,

571 + prefix=f"{prefix}.linear_fc2",

572 + disable_tp=use_data_parallel,

573 + )

574 + self.act_fn = act_fn

575 +

576 + def forward(self, x: torch.Tensor):

577 + mlp_output = self.linear_fc2(self.act_fn(self.linear_fc1(x)))

578 + return mlp_output

579 +

580 +

581 +class Cohere_VisionBlock(nn.Module):

582 + def __init__(

583 + self,

584 + dim: int,

585 + num_heads: int,

586 + mlp_hidden_dim: int,

587 + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

588 + norm_layer: Callable[[int], nn.Module] | None = None,

589 + quant_config: QuantizationConfig | None = None,

590 + prefix: str = "",

591 + ) -> None:

592 + super().__init__()

593 + if norm_layer is None:

594 + norm_layer = partial(nn.LayerNorm, eps=1e-6)

595 + self.norm1 = norm_layer(dim)

596 + self.norm2 = norm_layer(dim)

597 + self.attn = Cohere_VisionAttention(

598 + embed_dim=dim,

599 + num_heads=num_heads,

600 + projection_size=dim,

601 + quant_config=quant_config,

602 + prefix=f"{prefix}.attn",

603 + )

604 + self.mlp = Cohere_VisionMLP(

605 + dim,

606 + mlp_hidden_dim,

607 + act_fn=act_fn,

608 + bias=True,

609 + quant_config=quant_config,

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

611 + )

612 +

613 + def forward(

614 + self,

615 + x: torch.Tensor,

616 + cu_seqlens: torch.Tensor,

617 + rotary_pos_emb_cos: torch.Tensor,

618 + rotary_pos_emb_sin: torch.Tensor,

619 + max_seqlen: torch.Tensor, # Only used for Flash Attention

620 + sequence_lengths: torch.Tensor, # Only used for FlashInfer CuDNN backend

621 + ) -> torch.Tensor:

622 + x = x + self.attn(

623 + self.norm1(x),

624 + cu_seqlens=cu_seqlens,

625 + rotary_pos_emb_cos=rotary_pos_emb_cos,

626 + rotary_pos_emb_sin=rotary_pos_emb_sin,

627 + max_seqlen=max_seqlen,

628 + sequence_lengths=sequence_lengths,

629 + )

630 +

631 + x = x + self.mlp(self.norm2(x))

632 + return x

633 +

634 +

635 +class Cohere_VisionPatchMerger(nn.Module):

636 + def __init__(

637 + self,

638 + d_model: int,

639 + context_dim: int,

640 + norm_layer: Callable[[int], nn.Module] | None = None,

641 + spatial_merge_size: int = 2,

642 + use_postshuffle_norm: bool = False,

643 + quant_config: QuantizationConfig | None = None,

644 + prefix: str = "",

645 + ) -> None:

646 + super().__init__()

647 + use_data_parallel = is_vit_use_data_parallel()

648 + self.hidden_size = context_dim * (spatial_merge_size**2)

649 +

650 + self.use_postshuffle_norm = use_postshuffle_norm

651 + if self.use_postshuffle_norm:

652 + context_dim = self.hidden_size

653 +

654 + if norm_layer is None:

655 + norm_layer = partial(nn.LayerNorm, eps=1e-6)

656 + self.norm = norm_layer(context_dim)

657 + self.linear_fc1 = ColumnParallelLinear(

658 + self.hidden_size,

659 + self.hidden_size,

660 + bias=True,

661 + quant_config=quant_config,

662 + prefix=f"{prefix}.linear_fc1",

663 + disable_tp=use_data_parallel,

664 + )

665 + self.act_fn = nn.GELU()

666 + self.linear_fc2 = RowParallelLinear(

667 + self.hidden_size,

668 + d_model,

669 + bias=True,

670 + quant_config=quant_config,

671 + prefix=f"{prefix}.linear_fc2",

672 + disable_tp=use_data_parallel,

673 + )

674 +

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

676 + if self.use_postshuffle_norm:

677 + x = self.norm(x.view(-1, self.hidden_size))

678 + else:

679 + x = self.norm(x).view(-1, self.hidden_size)

680 +

681 + x_parallel, _ = self.linear_fc1(x)

682 + x_parallel = self.act_fn(x_parallel)

683 + out, _ = self.linear_fc2(x_parallel)

684 + return out

685 +

686 +

687 +class Cohere_VisionTransformer(nn.Module):

688 + hf_to_vllm_mapper = WeightsMapper(

689 + orig_to_new_stacked={

690 + ".q.": (".qkv.", "q"),

691 + ".k.": (".qkv.", "k"),

692 + ".v.": (".qkv.", "v"),

693 + }

694 + )

695 +

696 + def __init__(

697 + self,

698 + vision_config: CohereCompassVisionConfig,

699 + norm_eps: float = 1e-6,

700 + quant_config: QuantizationConfig | None = None,

701 + prefix: str = "",

702 + ) -> None:

703 + super().__init__()

704 + self.hidden_size = vision_config.hidden_size

705 + self.num_heads = vision_config.num_heads

706 + self.num_position_embeddings = vision_config.num_position_embeddings

707 + self.patch_size = vision_config.patch_size

708 + self.spatial_merge_size = vision_config.spatial_merge_size

709 + self.spatial_merge_unit = self.spatial_merge_size**2

710 + self.temporal_patch_size = vision_config.temporal_patch_size

711 + self.deepstack_visual_indexes = (

712 + vision_config.deepstack_visual_indexes

713 + if hasattr(vision_config, "deepstack_visual_indexes")

714 + else []

715 + )

716 + self.num_grid_per_side = int(self.num_position_embeddings**0.5)

717 +

718 + use_data_parallel = is_vit_use_data_parallel()

719 + self.tp_size = (

720 + 1

721 + if use_data_parallel

722 + else parallel_state.get_tensor_model_parallel_world_size()

723 + )

724 +

725 + # NOTE: This is used for creating empty tensor for all_gather for

726 + # DP ViT. Here out_hidden_size is enlarged due to deepstack

727 + self.out_hidden_size = vision_config.out_hidden_size * (

728 + 1 + len(self.deepstack_visual_indexes)

729 + )

730 +

731 + self.patch_embed = Cohere_VisionPatchEmbed(

732 + patch_size=self.patch_size,

733 + temporal_patch_size=self.temporal_patch_size,

734 + in_channels=vision_config.in_channels,

735 + hidden_size=self.hidden_size,

736 + )

737 +

738 + self.pos_embed = nn.Embedding(self.num_position_embeddings, self.hidden_size)

739 +

740 + norm_layer = partial(nn.LayerNorm, eps=norm_eps)

741 + head_dim = self.hidden_size // self.num_heads

742 +

743 + # FP8 attention: Q/K/V become independent contiguous tensors

744 + # after quantization, so cu_seqlens uses uniform stride (no 3x V).

745 + self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(

746 + self.num_heads, head_dim

747 + )

748 +

749 + self.rotary_pos_emb = get_rope(

750 + head_size=head_dim,

751 + max_position=8192,

752 + is_neox_style=True,

753 + rope_parameters={"partial_rotary_factor": 0.5},

754 + )

755 +

756 + self.merger = Cohere_VisionPatchMerger(

757 + d_model=vision_config.out_hidden_size,

758 + context_dim=self.hidden_size,

759 + norm_layer=norm_layer,

760 + spatial_merge_size=self.spatial_merge_size,

761 + quant_config=quant_config,

762 + prefix=f"{prefix}.merger",

763 + )

764 +

765 + self.deepstack_merger_list = nn.ModuleList(

766 + [

767 + Cohere_VisionPatchMerger(

768 + d_model=vision_config.out_hidden_size,

769 + context_dim=self.hidden_size,

770 + spatial_merge_size=self.spatial_merge_size,

771 + use_postshuffle_norm=True,

772 + norm_layer=norm_layer,

773 + quant_config=quant_config,

774 + prefix=f"{prefix}.deepstack_merger_list.{layer_idx}",

775 + )

776 + for layer_idx in range(len(self.deepstack_visual_indexes))

777 + ]

778 + )

779 +

780 + self.attn_backend = get_vit_attn_backend(

781 + head_size=head_dim,

782 + dtype=torch.get_default_dtype(),

783 + )

784 +

785 + self.blocks = nn.ModuleList(

786 + [

787 + Cohere_VisionBlock(

788 + dim=self.hidden_size,

789 + num_heads=self.num_heads,

790 + mlp_hidden_dim=vision_config.intermediate_size,

791 + act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],

792 + norm_layer=norm_layer,

793 + quant_config=quant_config,

794 + prefix=f"{prefix}.blocks.{layer_idx}",

795 + )

796 + for layer_idx in range(vision_config.depth)

797 + ]

798 + )

799 +

800 + @property

801 + def dtype(self) -> torch.dtype:

802 + return self.patch_embed.proj.weight.dtype

803 +

804 + @property

805 + def device(self) -> torch.device:

806 + return self.patch_embed.proj.weight.device

807 +

808 + @staticmethod

809 + @lru_cache(maxsize=1024)

810 + def rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor:

811 + hpos_ids = np.broadcast_to(np.arange(h).reshape(h, 1), (h, w))

812 + h_div = h // spatial_merge_size

813 + w_div = w // spatial_merge_size

814 + hpos_ids = hpos_ids.reshape(

815 + h_div,

816 + spatial_merge_size,

817 + w_div,

818 + spatial_merge_size,

819 + )

820 + hpos_ids = hpos_ids.transpose(0, 2, 1, 3)

821 + hpos_ids = hpos_ids.flatten()

822 +

823 + wpos_ids = np.broadcast_to(np.arange(w).reshape(1, w), (h, w))

824 + wpos_ids = wpos_ids.reshape(

825 + h_div,

826 + spatial_merge_size,

827 + w_div,

828 + spatial_merge_size,

829 + )

830 + wpos_ids = wpos_ids.transpose(0, 2, 1, 3)

831 + wpos_ids = wpos_ids.flatten()

832 +

833 + return torch.from_numpy(np.stack([hpos_ids, wpos_ids], axis=-1))

834 +

835 + def rot_pos_emb(self, grid_thw: list[list[int]]):

836 + max_grid_size = max(max(h, w) for _, h, w in grid_thw)

837 + pos_ids = [

838 + self.rot_pos_ids(h, w, self.spatial_merge_size)

839 + if t == 1

840 + else self.rot_pos_ids(h, w, self.spatial_merge_size).repeat(t, 1)

841 + for t, h, w in grid_thw

842 + ]

843 + pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True)

844 +

845 + # Use pre-computed cos_sin_cache from RotaryEmbedding

846 + cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

847 +

848 + cos_combined = cos[pos_ids].flatten(1)

849 + sin_combined = sin[pos_ids].flatten(1)

850 +

851 + return cos_combined, sin_combined

852 +

853 + def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:

854 + interpolate_fn = (

855 + triton_pos_embed_interpolate if HAS_TRITON else pos_embed_interpolate_native

856 + )

857 + outputs = []

858 + for t, h, w in grid_thw:

859 + outputs.append(

860 + interpolate_fn(

861 + self.pos_embed.weight,

862 + t,

863 + h,

864 + w,

865 + self.num_grid_per_side,

866 + self.spatial_merge_size,

867 + self.dtype,

868 + )

869 + )

870 + return torch.cat(outputs, dim=0)

871 +

872 + def prepare_encoder_metadata(

873 + self,

874 + grid_thw_list: list[list[int]],

875 + *,

876 + max_batch_size: int | None = None,

877 + max_frames_per_batch: int | None = None,

878 + max_seqlen_override: int | None = None,

879 + device: torch.device | None = None,

880 + ) -> dict[str, torch.Tensor | None]:

881 + """Compute encoder metadata from grid_thw_list.

882 +

883 + Shared by the eager forward path, CUDA graph capture, and

884 + CUDA graph replay to avoid duplicated implementation.

885 +

886 + Args:

887 + grid_thw_list: Grid configurations as list of [t, h, w].

888 + max_batch_size: If set, pad cu_seqlens to this size

889 + (needed for CUDA graph capture/replay).

890 + max_frames_per_batch: If set, overrides max_batch_size for

891 + cu_seqlens padding. For video inputs each item contributes

892 + T attention sequences (frames); this sizes the buffer to

893 + the total frame budget so video replays never overflow.

894 + max_seqlen_override: If set, use this value for max_seqlen

895 + instead of computing from cu_seqlens (needed for CUDA

896 + graph capture to cover worst-case replay scenarios).

897 + device: Device to place tensors on. Defaults to self.device.

898 + """

899 + if device is None:

900 + device = self.device

901 +

902 + metadata: dict[str, torch.Tensor | None] = {}

903 +

904 + # Positional embeddings

905 + metadata["pos_embeds"] = self.fast_pos_embed_interpolate(grid_thw_list)

906 + rotary_cos, rotary_sin = self.rot_pos_emb(grid_thw_list)

907 + metadata["rotary_pos_emb_cos"] = rotary_cos

908 + metadata["rotary_pos_emb_sin"] = rotary_sin

909 +

910 + # cu_seqlens from grid_thw

911 + grid_thw_np = np.array(grid_thw_list, dtype=np.int32)

912 + patches_per_frame = grid_thw_np[:, 1] * grid_thw_np[:, 2]

913 + cu_seqlens = np.repeat(patches_per_frame, grid_thw_np[:, 0]).cumsum(

914 + dtype=np.int32

915 + )

916 + cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])

917 +

918 + # Pad cu_seqlens to the required number of sequences.

919 + # For videos each item contributes T frames = T attention sequences,

920 + # so the total can exceed max_batch_size. max_frames_per_batch

921 + # overrides the pad target when set.

922 + pad_to = (

923 + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size

924 + )

925 + if pad_to is not None:

926 + num_seqs = len(cu_seqlens) - 1

927 + if num_seqs < pad_to:

928 + cu_seqlens = np.concatenate(

929 + [

930 + cu_seqlens,

931 + np.full(

932 + pad_to - num_seqs,

933 + cu_seqlens[-1],

934 + dtype=np.int32,

935 + ),

936 + ]

937 + )

938 +

939 + # sequence_lengths (backend-specific)

940 + metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(

941 + self.attn_backend, cu_seqlens, device

942 + )

943 +

944 + # max_seqlen

945 + if max_seqlen_override is not None:

946 + max_seqlen_val = max_seqlen_override

947 + else:

948 + max_seqlen_val = MMEncoderAttention.compute_max_seqlen(

949 + self.attn_backend, cu_seqlens

950 + )

951 + # Keep max_seqlen on CPU: attention wrappers call .item() on it,

952 + # and having it on GPU would capture a wasteful D2H copy in CUDA

953 + # graphs without changing behavior (the scalar is baked at capture).

954 + metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)

955 +

956 + # Recompute cu_seqlens (backend-specific transformation)

957 + metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(

958 + self.attn_backend,

959 + cu_seqlens,

960 + self.hidden_size,

961 + self.tp_size,

962 + device,

963 + fp8_padded_hidden_size=self.fp8_padded_hidden_size,

964 + )

965 +

966 + return metadata

967 +

968 + def forward(

969 + self,

970 + x: torch.Tensor,

971 + grid_thw: torch.Tensor | list[list[int]],

972 + *,

973 + encoder_metadata: dict[str, torch.Tensor] | None = None,

974 + ) -> torch.Tensor:

975 + hidden_states = x.to(device=self.device, dtype=self.dtype, non_blocking=True)

976 + hidden_states = self.patch_embed(hidden_states)

977 +

978 + if encoder_metadata is None:

979 + if isinstance(grid_thw, list):

980 + grid_thw_list = grid_thw

981 + else:

982 + grid_thw_list = grid_thw.tolist()

983 + encoder_metadata = self.prepare_encoder_metadata(grid_thw_list)

984 +

985 + pos_embeds = encoder_metadata["pos_embeds"]

986 + hidden_states = hidden_states + pos_embeds

987 + hidden_states = hidden_states.unsqueeze(1)

988 +

989 + deepstack_feature_lists = []

990 + for layer_num, blk in enumerate(self.blocks):

991 + hidden_states = blk(

992 + hidden_states,

993 + cu_seqlens=encoder_metadata["cu_seqlens"],

994 + rotary_pos_emb_cos=encoder_metadata["rotary_pos_emb_cos"],

995 + rotary_pos_emb_sin=encoder_metadata["rotary_pos_emb_sin"],

996 + max_seqlen=encoder_metadata["max_seqlen"],

997 + sequence_lengths=encoder_metadata.get("sequence_lengths"),

998 + )

999 + if layer_num in self.deepstack_visual_indexes:

1000 + deepstack_merger_idx = self.deepstack_visual_indexes.index(layer_num)

1001 + deepstack_feature = self.deepstack_merger_list[deepstack_merger_idx](

1002 + hidden_states

1003 + )

1004 + deepstack_feature_lists.append(deepstack_feature)

1005 + hidden_states = self.merger(hidden_states)

1006 + hidden_states = torch.cat(

1007 + [hidden_states] + deepstack_feature_lists, dim=1

1008 + ) # [seq_len, hidden_size * (1 + depth_of_deepstack)]

1009 + return hidden_states

1010 +

1011 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1012 + loader = AutoWeightsLoader(self)

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

1014 +

1015 +

1016 +def _create_cohere_compass_field_factory(

1017 + spatial_merge_size: int,

1018 +) -> Callable[

1019 + [Mapping[str, torch.Tensor]],

1020 + Mapping[str, MultiModalFieldConfig],

1021 +]:

1022 + def _field_config(

1023 + hf_inputs: Mapping[str, torch.Tensor],

1024 + ) -> Mapping[str, MultiModalFieldConfig]:

1025 + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))

1026 + image_pixel_grid_sizes = image_grid_thw.prod(-1)

1027 + image_embed_grid_sizes = (

1028 + image_pixel_grid_sizes // spatial_merge_size // spatial_merge_size

1029 + )

1030 +

1031 + return {

1032 + "pixel_values": MultiModalFieldConfig.flat_from_sizes(

1033 + "image", image_pixel_grid_sizes

1034 + ),

1035 + "image_embeds": MultiModalFieldConfig.flat_from_sizes(

1036 + "image", image_embed_grid_sizes

1037 + ),

1038 + "image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True),

1039 + }

1040 +

1041 + return _field_config

1042 +

1043 +

1044 +class CohereCompassMultiModalDataParser(MultiModalDataParser):

1045 + def __init__(self, spatial_merge_size: int, *args, **kwargs):

1046 + self._spatial_merge_size = spatial_merge_size

1047 + super().__init__(*args, **kwargs)

1048 +

1049 + def _parse_image_data(

1050 + self,

1051 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

1052 + ) -> ModalityDataItems[Any, Any] | None:

1053 + if isinstance(data, dict):

1054 + return DictEmbeddingItems(

1055 + data,

1056 + modality="image",

1057 + required_fields={"image_embeds", "image_grid_thw"},

1058 + fields_factory=_create_cohere_compass_field_factory(

1059 + self._spatial_merge_size

1060 + ),

1061 + )

1062 +

1063 + return super()._parse_image_data(data)

1064 +

1065 +

1066 +class CohereCompassProcessingInfo(BaseProcessingInfo):

1067 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

1068 + return {"image": None}

1069 +

1070 + def get_mm_max_tokens_per_item(

1071 + self,

1072 + seq_len: int,

1073 + mm_counts: Mapping[str, int],

1074 + ) -> Mapping[str, int]:

1075 + max_image_tokens = self.get_max_image_tokens()

1076 + return {"image": max_image_tokens}

1077 +

1078 + def get_num_image_tokens(

1079 + self,

1080 + *,

1081 + image_width: int,

1082 + image_height: int,

1083 + image_processor: CohereCompassImageProcessor,

1084 + mm_kwargs: Mapping[str, object],

1085 + ) -> int:

1086 + _, num_image_tokens = self._get_vision_info(

1087 + image_width=image_width,

1088 + image_height=image_height,

1089 + num_frames=1,

1090 + image_processor=image_processor,

1091 + mm_kwargs=mm_kwargs,

1092 + )

1093 + return num_image_tokens

1094 +

1095 + def get_image_size_with_most_features(

1096 + self, max_pixels: int | None = None

1097 + ) -> ImageSize:

1098 + # NOTE: Simply processing a huge size with _get_vision_info might not give a

1099 + # size that maximizes the number of features, i.e., the number of (merged)

1100 + # patches. This is because the number of patches limits the allowed aspect

1101 + # ratios. For example, suppose the maximum number of patches is 1280. A square

1102 + # image cannot be broken down into 1280 patches, so feeding a giant square image

1103 + # into _get_vision_info will not yield a size that maximizes the number of

1104 + # patches. Therefore, we directly factorize the maximum number of patches into

1105 + # height and width. The tricky part is to avoid extreme aspect ratios (>200).

1106 + # If we can't find a suitable aspect ratio, we decrease the number of

1107 + # patches and retry. This is safe because the processor does not accept extreme

1108 + # aspect ratios, so there is no valid post-resize image with the number of

1109 + # patches that yields extreme aspect ratios.

1110 +

1111 + hf_config = self.get_hf_config()

1112 + vision_config = hf_config.vision_config

1113 + patch_size = vision_config.patch_size

1114 + merge_size = vision_config.spatial_merge_size

1115 +

1116 + if max_pixels is None:

1117 + image_processor = self.get_image_processor()

1118 +

1119 + mm_kwargs = self.ctx.get_merged_mm_kwargs({})

1120 + size = image_processor.size

1121 + if override_size := mm_kwargs.get("size"):

1122 + size = size | override_size

1123 + if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None:

1124 + size = size | {"shortest_edge": override_min_pixels}

1125 + if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None:

1126 + size = size | {"longest_edge": override_max_pixels}

1127 +

1128 + max_pixels = size["longest_edge"]

1129 +

1130 + unit = patch_size * merge_size

1131 + max_seq_len = max_pixels // (unit * unit)

1132 +

1133 + def closest_factor_pair(n: int) -> tuple[int, int]:

1134 + # left <= right

1135 + for d in range(math.isqrt(n), 0, -1):

1136 + if n % d == 0:

1137 + return d, n // d

1138 + return 1, n

1139 +

1140 + height_factor, width_factor = 1, max_seq_len

1141 + for seq_len in range(max_seq_len, 0, -1):

1142 + height_factor, width_factor = closest_factor_pair(seq_len)

1143 + if width_factor / height_factor <= 200:

1144 + break

1145 +

1146 + return ImageSize(width=unit * width_factor, height=unit * height_factor)

1147 +

1148 + def get_max_image_tokens(self) -> int:

1149 + image_processor = self.get_image_processor()

1150 + target_width, target_height = self.get_image_size_with_most_features()

1151 +

1152 + return self.get_num_image_tokens(

1153 + image_width=target_width,

1154 + image_height=target_height,

1155 + image_processor=image_processor,

1156 + mm_kwargs={},

1157 + )

1158 +

1159 + def get_hf_config(self) -> CohereCompassConfig:

1160 + return self.ctx.get_hf_config(CohereCompassConfig)

1161 +

1162 + def get_hf_processor(self, **kwargs: object) -> CohereCompassProcessor:

1163 + return self.ctx.get_hf_processor(

1164 + CohereCompassProcessor,

1165 + use_fast=kwargs.pop("use_fast", True),

1166 + **kwargs,

1167 + )

1168 +

1169 + def get_image_processor(self, **kwargs: object) -> CohereCompassImageProcessor:

1170 + return self.get_hf_processor(**kwargs).image_processor

1171 +

1172 + def get_data_parser(self):

1173 + return CohereCompassMultiModalDataParser(

1174 + self.get_hf_config().vision_config.spatial_merge_size,

1175 + expected_hidden_size=self._get_expected_hidden_size(),

1176 + )

1177 +

1178 + def _get_vision_info(

1179 + self,

1180 + *,

1181 + image_width: int,

1182 + image_height: int,

1183 + num_frames: int = 2,

1184 + do_resize: bool = True,

1185 + image_processor: CohereCompassImageProcessor,

1186 + mm_kwargs: Mapping[str, object],

1187 + ) -> tuple[ImageSize, int]:

1188 + hf_config = self.get_hf_config()

1189 + vision_config = hf_config.vision_config

1190 + patch_size = vision_config.patch_size

1191 + merge_size = vision_config.spatial_merge_size

1192 + temporal_patch_size = vision_config.temporal_patch_size

1193 +

1194 + mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs)

1195 + size = image_processor.size

1196 + if override_size := mm_kwargs.get("size"):

1197 + size = size | override_size

1198 + if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None:

1199 + size = size | {"shortest_edge": override_min_pixels}

1200 + if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None:

1201 + size = size | {"longest_edge": override_max_pixels}

1202 +

1203 + if do_resize:

1204 + smart_resize = image_smart_resize

1205 + extra_kwargs: dict[str, Any] = {}

1206 +

1207 + resized_height, resized_width = smart_resize(

1208 + height=image_height,

1209 + width=image_width,

1210 + factor=patch_size * merge_size,

1211 + min_pixels=size["shortest_edge"],

1212 + max_pixels=size["longest_edge"],

1213 + **extra_kwargs,

1214 + )

1215 + preprocessed_size = ImageSize(width=resized_width, height=resized_height)

1216 + else:

1217 + preprocessed_size = ImageSize(width=image_width, height=image_height)

1218 +

1219 + padded_num_frames = round_up(num_frames, temporal_patch_size)

1220 +

1221 + grid_t = max(padded_num_frames // temporal_patch_size, 1)

1222 + grid_h = preprocessed_size.height // patch_size

1223 + grid_w = preprocessed_size.width // patch_size

1224 +

1225 + num_patches = grid_t * grid_h * grid_w

1226 + num_vision_tokens = num_patches // (merge_size**2)

1227 +

1228 + return preprocessed_size, num_vision_tokens

1229 +

1230 +

1231 +class CohereCompassDummyInputsBuilder(

1232 + BaseDummyInputsBuilder[CohereCompassProcessingInfo]

1233 +):

1234 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

1235 + num_images = mm_counts.get("image", 0)

1237 +

1238 + def get_dummy_mm_data(

1239 + self,

1240 + seq_len: int,

1241 + mm_counts: Mapping[str, int],

1242 + mm_options: Mapping[str, BaseDummyOptions],

1243 + ) -> MultiModalDataDict:

1244 + num_images = mm_counts.get("image", 0)

1245 + image_overrides = mm_options.get("image")

1246 +

1247 + target_image_width, target_image_height = (

1248 + self.info.get_image_size_with_most_features()

1249 + )

1250 +

1251 + return {

1252 + "image": self._get_dummy_images(

1253 + width=target_image_width,

1254 + height=target_image_height,

1255 + num_images=num_images,

1256 + overrides=image_overrides,

1257 + ),

1258 + }

1259 +

1260 +

1261 +class CohereCompassMultiModalProcessor(

1262 + BaseMultiModalProcessor[CohereCompassProcessingInfo]

1263 +):

1264 + def _get_mm_fields_config(

1265 + self,

1266 + hf_inputs: BatchFeature,

1267 + hf_processor_mm_kwargs: Mapping[str, object],

1268 + ) -> Mapping[str, MultiModalFieldConfig]:

1269 + return _create_cohere_compass_field_factory(

1270 + self.info.get_hf_config().vision_config.spatial_merge_size

1271 + )(hf_inputs)

1272 +

1273 + def _get_prompt_updates(

1274 + self,

1275 + mm_items: MultiModalDataItems,

1276 + hf_processor_mm_kwargs: Mapping[str, Any],

1277 + out_mm_kwargs: MultiModalKwargsItems,

1278 + ) -> Sequence[PromptUpdate]:

1279 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

1280 + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)

1281 + merge_length = image_processor.merge_size**2

1282 +

1283 + def get_image_replacement_cohere_compass(item_idx: int):

1284 + out_item = out_mm_kwargs["image"][item_idx]

1285 + grid_thw = out_item["image_grid_thw"].data

1286 + assert isinstance(grid_thw, torch.Tensor)

1287 +

1288 + num_tokens = int(grid_thw.prod()) // merge_length

1289 + return [hf_processor.image_token_id] * num_tokens

1290 +

1291 + return [

1292 + PromptReplacement(

1293 + modality="image",

1294 + target=[hf_processor.image_token_id],

1295 + replacement=get_image_replacement_cohere_compass,

1296 + ),

1297 + ]

1298 +

1299 +

1300 +@MULTIMODAL_REGISTRY.register_processor(

1301 + CohereCompassMultiModalProcessor,

1302 + info=CohereCompassProcessingInfo,

1303 + dummy_inputs=CohereCompassDummyInputsBuilder,

1304 +)

1305 +class CohereCompassForConditionalGeneration(

1306 + nn.Module,

1307 + SupportsMultiModal,

1308 + SupportsEncoderCudaGraph,

1309 + SupportsLoRA,

1310 + SupportsPP,

1311 + SupportsMRoPE,

1312 + SupportsEagle,

1313 + SupportsEagle3,

1314 + SupportsMultiModalPruning,

1315 +):

1316 + packed_modules_mapping = {

1317 + "qkv_proj": [

1318 + "q_proj",

1319 + "k_proj",

1320 + "v_proj",

1321 + ],

1322 + "gate_up_proj": [

1323 + "gate_proj",

1324 + "up_proj",

1325 + ],

1326 + "qkv": ["qkv"], # For vision tower's already-packed QKV

1327 + }

1328 +

1329 + supports_encoder_tp_data = True

1330 +

1331 + # To ensure correct weight loading and mapping.

1332 + hf_to_vllm_mapper = WeightsMapper(

1333 + orig_to_new_prefix={

1334 + "model.visual.": "visual.",

1335 + "lm_head.": "language_model.lm_head.",

1336 + "model.language_model.": "language_model.model.",

1337 + }

1338 + )

1339 +

1340 + @classmethod

1341 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1342 + if modality.startswith("image"):

1343 + return "<|VISION_START|><|IMAGE_PAD|><|VISION_END|>"

1344 + raise ValueError("Only image modality is supported")

1345 +

1346 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):

1347 + nn.Module.__init__(self)

1348 + config = vllm_config.model_config.hf_config

1349 + quant_config = vllm_config.quant_config

1350 + multimodal_config = vllm_config.model_config.multimodal_config

1351 +

1352 + self.config = config

1353 + self.model_config = vllm_config.model_config

1354 + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)

1355 + self.multimodal_config = multimodal_config

1356 + assert multimodal_config is not None

1357 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1358 + self.is_multimodal_pruning_enabled = (

1359 + multimodal_config.is_multimodal_pruning_enabled()

1360 + )

1361 +

1362 + self.use_deepstack = bool(config.vision_config.deepstack_visual_indexes)

1363 + self.deepstack_num_level = len(config.vision_config.deepstack_visual_indexes)

1364 + self.visual_dim = config.vision_config.out_hidden_size

1365 + self.multiscale_dim = self.visual_dim * self.deepstack_num_level

1366 +

1367 + with self._mark_tower_model(vllm_config, {"image"}):

1368 + self.visual = Cohere_VisionTransformer(

1369 + config.vision_config,

1370 + norm_eps=1e-6,

1371 + quant_config=quant_config,

1372 + prefix=maybe_prefix(prefix, "visual"),

1373 + )

1374 + if self.use_deepstack:

1375 + self.deepstack_input_embeds = [

1376 + torch.zeros(

1377 + vllm_config.scheduler_config.max_num_batched_tokens,

1378 + config.text_config.hidden_size,

1379 + )

1380 + for _ in range(self.deepstack_num_level)

1381 + ]

1382 + self.deepstack_input_embeds_num_tokens = 0

1383 +

1384 + with self._mark_language_model(vllm_config):

1385 + self.language_model = CohereCompassForCausalLM(

1386 + vllm_config=vllm_config.with_hf_config(config.text_config),

1387 + prefix=maybe_prefix(prefix, "language_model"),

1388 + )

1389 +

1390 + if not get_pp_group().is_first_rank and self.use_deepstack:

1391 + assert self.language_model.model.start_layer >= self.deepstack_num_level

1392 +

1393 + self.make_empty_intermediate_tensors = (

1394 + self.language_model.make_empty_intermediate_tensors

1395 + )

1396 +

1397 + def iter_mm_grid_thw(

1398 + self, mm_features: list[MultiModalFeatureSpec]

1399 + ) -> Iterator[tuple[int, int, int, int, float]]:

1400 + spatial_merge_size = self.config.vision_config.spatial_merge_size

1401 + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

1402 + if mm_feature.modality != "image":

1403 + raise ValueError(f"Unsupported modality: {mm_feature.modality}")

1404 +

1405 + offset = mm_feature.mm_position.offset

1406 + feature_data = mm_feature.data

1407 + assert feature_data is not None

1408 + grid_item = feature_data.get("image_grid_thw")

1409 + assert grid_item is not None

1410 + grid_data = grid_item.data

1411 + assert isinstance(grid_data, torch.Tensor)

1412 + t, h, w = grid_data.tolist()

1413 + assert t == 1, f"Image must have 1 frame, got {t}"

1414 + yield offset, 1, h // spatial_merge_size, w // spatial_merge_size, 1.0

1415 +

1416 + def _get_deepstack_input_embeds(

1417 + self,

1418 + num_tokens: int,

1419 + ) -> IntermediateTensors | None:

1420 + if not getattr(self, "deepstack_input_embeds", None):

1421 + return None # If vision tower is skipped

1422 + if num_tokens > self.deepstack_input_embeds[0].size(0):

1423 + self._resize_deepstack_input_embeds(num_tokens)

1424 +

1425 + # get deepstack_input_embeds from buffer, and clear the buffer

1426 + return IntermediateTensors(

1427 + {

1428 + f"deepstack_input_embeds_{idx}": self.deepstack_input_embeds[idx][

1429 + :num_tokens

1430 + ]

1431 + for idx in range(self.deepstack_num_level)

1432 + }

1433 + )

1434 +

1435 + def _resize_deepstack_input_embeds(self, num_tokens: int) -> None:

1436 + self.deepstack_input_embeds = [

1437 + torch.zeros(

1438 + num_tokens,

1439 + self.config.text_config.hidden_size,

1440 + device=self.deepstack_input_embeds[0].device,

1441 + dtype=self.deepstack_input_embeds[0].dtype,

1442 + )

1443 + for _ in range(self.deepstack_num_level)

1444 + ]

1445 +

1446 + def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None:

1447 + if not getattr(self, "deepstack_input_embeds", None):

1448 + return

1449 +

1450 + # set deepstack_input_embeds to buffer

1451 + num_tokens = deepstack_input_embeds.size(1)

1452 + if num_tokens > self.deepstack_input_embeds[0].size(0):

1453 + self._resize_deepstack_input_embeds(num_tokens)

1454 + for idx in range(self.deepstack_num_level):

1455 + self.deepstack_input_embeds[idx][:num_tokens].copy_(

1456 + deepstack_input_embeds[idx]

1457 + )

1458 + self.deepstack_input_embeds_num_tokens = num_tokens

1459 +

1460 + def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:

1461 + if not getattr(self, "deepstack_input_embeds", None):

1462 + return

1463 + if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:

1464 + return

1465 +

1466 + # clear deepstack_input_embeds in buffer

1467 + if num_tokens > 0:

1468 + for idx in range(self.deepstack_num_level):

1469 + self.deepstack_input_embeds[idx][:num_tokens].zero_()

1470 + self.deepstack_input_embeds_num_tokens = 0

1471 +

1472 + # -- SupportsEncoderCudaGraph protocol methods --

1473 +

1474 + def get_encoder_cudagraph_config(self):

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

1476 + EncoderCudaGraphConfig,

1477 + )

1478 +

1479 + # When EVS pruning is enabled, embed_multimodal post-processes both

1480 + # image and video embeddings (mrope positions are appended for image,

1481 + # prune+append for video). The encoder CUDA graph path bypasses that

1482 + # post-process, producing inconsistent embedding formats vs eager. So

1483 + # disable CUDA graph for all modalities when pruning is on.

1484 + modalities = [] if self.is_multimodal_pruning_enabled else ["image"]

1485 +

1486 + # Compute max_frames_per_video for budget sizing.

1487 + max_frames = 1

1488 +

1489 + return EncoderCudaGraphConfig(

1490 + modalities=modalities,

1491 + buffer_keys=[

1492 + "pixel_values",

1493 + "pos_embeds",

1494 + "rotary_pos_emb_cos",

1495 + "rotary_pos_emb_sin",

1496 + "cu_seqlens",

1497 + "max_seqlen",

1498 + "sequence_lengths",

1499 + ],

1500 + out_hidden_size=self.visual.out_hidden_size,

1501 + max_frames_per_video=max_frames,

1502 + )

1503 +

1504 + def get_input_modality(

1505 + self,

1506 + mm_kwargs: dict[str, Any],

1507 + ) -> str:

1508 + if "image_grid_thw" in mm_kwargs:

1509 + return "image"

1510 + raise AssertionError("This line should be unreachable.")

1511 +

1512 + def get_encoder_cudagraph_budget_range(

1513 + self,

1514 + vllm_config,

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

1516 + # Min: estimated smallest possible encoder input.

1517 + # 224x224 image → 16x16 patches (patch_size=14)

1518 + # spatial_merge_size=2 → 8x8 = 64 tokens

1519 + min_budget = 64

1520 + # Max: capped by max_num_batched_tokens

1521 + # TODO(shen-shanshan): the max_budget auto-infer needs to be optimized later.

1522 + max_budget = min(

1523 + vllm_config.scheduler_config.max_num_batched_tokens,

1524 + self.model_config.max_model_len,

1525 + )

1526 + return (min_budget, max_budget)

1527 +

1528 + def _get_pixel_values_by_modality(

1529 + self,

1530 + mm_kwargs: dict[str, Any],

1531 + ) -> torch.Tensor:

1532 + modality = self.get_input_modality(mm_kwargs)

1533 + if modality == "image":

1534 + return mm_kwargs["pixel_values"]

1535 + raise AssertionError("This line should be unreachable.")

1536 +

1537 + def _get_grid_thw_by_modality(

1538 + self,

1539 + mm_kwargs: dict[str, Any],

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

1541 + grid_thw_key = f"{self.get_input_modality(mm_kwargs)}_grid_thw"

1542 + grid_thw = mm_kwargs[grid_thw_key]

1543 + if not isinstance(grid_thw, list):

1544 + grid_thw = grid_thw.tolist()

1545 + return grid_thw

1546 +

1547 + def get_encoder_cudagraph_item_specs(

1548 + self,

1549 + mm_kwargs: dict[str, Any],

1550 + ):

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

1552 +

1553 + m = self.visual.spatial_merge_size

1554 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1555 + return [

1556 + EncoderItemSpec(

1557 + input_size=t * h * w,

1558 + output_tokens=t * (h // m) * (w // m),

1559 + )

1560 + for t, h, w in grid_thw

1561 + ]

1562 +

1563 + def select_encoder_cudagraph_items(

1564 + self,

1565 + mm_kwargs: dict[str, Any],

1566 + indices: list[int],

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

1568 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1569 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1570 +

1571 + if len(indices) == 0:

1572 + return {

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

1574 + "image_grid_thw": [],

1575 + }

1576 +

1577 + # Compute cumulative patch offsets for slicing pixel_values

1578 + patches_per_item = [t * h * w for t, h, w in grid_thw]

1579 + cum_patches = [0]

1580 + for p in patches_per_item:

1581 + cum_patches.append(cum_patches[-1] + p)

1582 +

1583 + selected_pv = torch.cat(

1584 + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]

1585 + )

1586 + selected_grid = [grid_thw[i] for i in indices]

1587 +

1588 + return {

1589 + "pixel_values": selected_pv,

1590 + "image_grid_thw": selected_grid,

1591 + }

1592 +

1593 + def prepare_encoder_cudagraph_capture_inputs(

1594 + self,

1595 + token_budget: int,

1596 + max_batch_size: int,

1597 + max_frames_per_batch: int,

1598 + device: torch.device,

1599 + dtype: torch.dtype,

1600 + path: str = "default",

1601 + ):

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

1603 + EncoderCudaGraphCaptureInputs,

1604 + )

1605 +

1606 + spatial_merge_size = self.visual.spatial_merge_size

1607 + # Ceil so the buffer fits the worst case of one item using the full

1608 + # budget. Floor under-allocates when budget is not a multiple of

1609 + # max_batch_size.

1610 + per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size

1611 +

1612 + # Image-format grid_config (T=1).

1613 + grid_config = [

1614 + [1, spatial_merge_size, per_mm_item_output * spatial_merge_size]

1615 + for _ in range(max_batch_size)

1616 + ]

1617 +

1618 + # Create dummy pixel_values

1619 + patch_embed = self.visual.patch_embed

1620 + in_channels = patch_embed.proj.in_channels

1621 + patch_size = patch_embed.patch_size

1622 + temporal_patch_size = patch_embed.temporal_patch_size

1623 + total_patches = sum(t * h * w for t, h, w in grid_config)

1624 + flattened_patch_size = (

1625 + in_channels * temporal_patch_size * patch_size * patch_size

1626 + )

1627 + dummy_pixel_values = torch.randn(

1628 + total_patches, flattened_patch_size, device=device, dtype=dtype

1629 + )

1630 +

1631 + # Override max_seqlen with a safe upper bound for capture.

1632 + # max_seqlen.item() gets baked into the CUDA graph (not replayed),

1633 + # so the capture value must cover any replay scenario.

1634 + # Worst case: 1 item consuming the full budget ->

1635 + # seq_len = token_budget * spatial_merge_size^2.

1636 + metadata = self.visual.prepare_encoder_metadata(

1637 + grid_config,

1638 + max_batch_size=max_batch_size,

1639 + max_frames_per_batch=max_frames_per_batch,

1640 + max_seqlen_override=token_budget * (spatial_merge_size**2),

1641 + device=device,

1642 + )

1643 +

1644 + # Just use image-modality dummy input_buffer for capturing

1645 + values = metadata | {

1646 + "pixel_values": dummy_pixel_values,

1647 + }

1648 +

1649 + return EncoderCudaGraphCaptureInputs(

1650 + values=values,

1651 + )

1652 +

1653 + def prepare_encoder_cudagraph_replay_buffers(

1654 + self,

1655 + mm_kwargs: dict[str, Any],

1656 + max_batch_size: int,

1657 + max_frames_per_batch: int,

1658 + path: str = "default",

1659 + ) -> EncoderCudaGraphReplayBuffers:

1660 + modality = self.get_input_modality(mm_kwargs)

1661 + grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)

1662 +

1663 + if modality == "image":

1664 + metadata = self.visual.prepare_encoder_metadata(

1665 + grid_thw_list,

1666 + max_batch_size=max_batch_size,

1667 + )

1668 + else:

1669 + raise AssertionError("This line should be unreachable.")

1670 +

1671 + values = metadata | {

1672 + "pixel_values": self._get_pixel_values_by_modality(mm_kwargs),

1673 + }

1674 + return EncoderCudaGraphReplayBuffers(values=values)

1675 +

1676 + def encoder_cudagraph_forward(

1677 + self,

1678 + values: dict[str, torch.Tensor],

1679 + path: str = "default",

1680 + ) -> torch.Tensor:

1681 + pixel_values = values.pop("pixel_values")

1682 + metadata = values

1683 + return self.visual(pixel_values, None, encoder_metadata=metadata)

1684 +

1685 + def encoder_eager_forward(

1686 + self,

1687 + mm_kwargs: dict[str, Any],

1688 + path: str = "default",

1689 + ) -> torch.Tensor:

1690 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1691 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1692 + return self.visual(pixel_values, grid_thw)

1693 +

1694 + def _parse_and_validate_image_input(

1695 + self, **kwargs: object

1696 + ) -> Cohere_VLImageInputs | None:

1697 + pixel_values = kwargs.pop("pixel_values", None)

1698 + image_embeds = kwargs.pop("image_embeds", None)

1699 + image_grid_thw = kwargs.pop("image_grid_thw", None)

1700 +

1701 + if pixel_values is None and image_embeds is None:

1702 + return None

1703 +

1704 + if pixel_values is not None:

1705 + return Cohere_VLImagePixelInputs(

1706 + type="pixel_values",

1707 + pixel_values=pixel_values,

1708 + image_grid_thw=image_grid_thw,

1709 + )

1710 +

1711 + if image_embeds is not None:

1712 + return Cohere_VLImageEmbeddingInputs(

1713 + type="image_embeds",

1714 + image_embeds=image_embeds,

1715 + image_grid_thw=image_grid_thw,

1716 + )

1717 + raise AssertionError("image input must contain pixels or embeddings")

1718 +

1719 + def _process_image_input(

1720 + self, image_input: Cohere_VLImageInputs

1721 + ) -> tuple[torch.Tensor, ...]:

1722 + grid_thw = image_input["image_grid_thw"]

1723 + assert grid_thw.ndim == 2

1724 +

1725 + if image_input["type"] == "image_embeds":

1726 + image_embeds = image_input["image_embeds"].type(self.visual.dtype)

1727 + else:

1728 + pixel_values = image_input["pixel_values"].type(self.visual.dtype)

1729 + if self.use_data_parallel:

1730 + return run_dp_sharded_mrope_vision_model(

1731 + self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"

1732 + )

1733 + else:

1734 + image_embeds = self.visual(pixel_values, grid_thw=grid_thw)

1735 +

1736 + # Split concatenated embeddings for each image item.

1737 + merge_size = self.visual.spatial_merge_size

1738 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1739 + return image_embeds.split(sizes)

1740 +

1741 + def _postprocess_image_embeds_evs(

1742 + self,

1743 + image_embeds_split: tuple[torch.Tensor, ...],

1744 + image_input: Cohere_VLImageInputs,

1745 + ) -> tuple[torch.Tensor, ...]:

1746 + """

1747 + Append mrope positions for each for images.

1748 + This is necessary to recover correct mrope positions

1749 +

1750 + Args:

1751 + image_embeds_split: Tuple of image embeddings for

1752 + each image item.

1753 + image_input: Image input data.

1754 +

1755 + Returns:

1756 + Tuple of image embeddings for each image item.

1757 + Resulting embeddings will have extra 5 channels for

1758 + computed mrope positions, consistent with video embeddings.

1759 + """

1760 + if self.is_multimodal_pruning_enabled:

1761 + merge_size = self.visual.spatial_merge_size

1762 + grid_thw = image_input["image_grid_thw"]

1763 + grid_thw_list = grid_thw.tolist()

1764 + image_embeds_out = []

1765 + for emb, size in zip(image_embeds_split, grid_thw_list):

1766 + positions = compute_mrope_for_media(size, merge_size).to(

1767 + emb.device, non_blocking=True

1768 + )

1769 + positions = torch.cat(

1770 + [

1771 + positions,

1772 + torch.zeros_like(

1773 + positions[:, 0:1]

1774 + ), # Dummy extra fifth channel

1775 + ],

1776 + dim=1,

1777 + )

1778 + emb = torch.cat([emb, positions], dim=1)

1779 + image_embeds_out.append(emb)

1780 + image_embeds_split = tuple(image_embeds_out)

1781 + return image_embeds_split

1782 +

1783 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1784 + mm_input_by_modality = {}

1785 + for input_key in kwargs:

1786 + if (

1787 + input_key in ("pixel_values", "image_embeds")

1788 + and "image" not in mm_input_by_modality

1789 + ):

1790 + mm_input_by_modality["image"] = self._parse_and_validate_image_input(

1791 + **kwargs

1792 + )

1793 + return mm_input_by_modality

1794 +

1795 + def get_mrope_input_positions(

1796 + self,

1797 + input_tokens: list[int],

1798 + mm_features: list[MultiModalFeatureSpec],

1799 + ) -> tuple[torch.Tensor, int]:

1800 + llm_pos_ids_list: list[np.ndarray] = []

1801 + st = 0

1802 +

1803 + for (

1804 + offset,

1805 + llm_grid_t,

1806 + llm_grid_h,

1807 + llm_grid_w,

1808 + _,

1809 + ) in self.iter_mm_grid_thw(mm_features):

1810 + text_len = offset - st

1811 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0

1812 + llm_pos_ids_list.append(

1813 + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx

1814 + )

1815 + grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w)).reshape(

1816 + 3, -1

1817 + )

1818 + llm_pos_ids_list.append(grid_indices + text_len + st_idx)

1819 + st = offset + llm_grid_t * llm_grid_h * llm_grid_w

1820 +

1821 + if st < len(input_tokens):

1822 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0

1823 + text_len = len(input_tokens) - st

1824 + llm_pos_ids_list.append(

1825 + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx

1826 + )

1827 +

1828 + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)

1829 + mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()

1830 + return torch.from_numpy(llm_positions), mrope_position_delta

1831 +

1832 + def recompute_mrope_positions(

1833 + self,

1834 + input_ids: list[int],

1835 + multimodal_embeddings: MultiModalEmbeddings,

1836 + mrope_positions: torch.LongTensor,

1837 + num_computed_tokens: int,

1838 + ) -> tuple[MultiModalEmbeddings, torch.Tensor, int]:

1839 + """

1840 + Update part of input mrope positions (starting with

1841 + num_computed_tokens index). Original mrope_positions are computed

1842 + for unpruned sequence and becomes incorrect once pruning occurs,

1843 + so once we prune media tokens we should reflect this in the

1844 + mrope_positions before we feed it to LLM.

1845 +

1846 + Args:

1847 + input_ids: (N,) All input tokens of the prompt containing

1848 + entire sequence.

1849 + multimodal_embeddings: Tuple of multimodal embeddings that

1850 + fits into the prefill chunk that is being processed.

1851 + mrope_positions: Existing mrope positions (3, N) for entire

1852 + sequence

1853 + num_computed_tokens: A number of computed tokens so far.

1854 +

1855 + Returns:

1856 + Tuple of (multimodal_embeddings, mrope_positions,

1857 + mrope_position_delta).

1858 + """

1859 + return self._recompute_mrope_positions(

1860 + input_ids=input_ids,

1861 + multimodal_embeddings=multimodal_embeddings,

1862 + mrope_positions=mrope_positions,

1863 + num_computed_tokens=num_computed_tokens,

1864 + image_token_id=self.config.image_token_id,

1865 + video_token_id=self.config.video_token_id,

1866 + vision_start_token_id=self.config.vision_start_token_id,

1867 + )

1868 +

1869 + @staticmethod

1870 + def _recompute_mrope_positions(

1871 + input_ids: list[int],

1872 + multimodal_embeddings: MultiModalEmbeddings,

1873 + mrope_positions: torch.LongTensor,

1874 + num_computed_tokens: int,

1875 + vision_start_token_id: int,

1876 + image_token_id: int,

1877 + video_token_id: int,

1878 + ) -> tuple[MultiModalEmbeddings, torch.Tensor, int]:

1879 + # Device

1880 + device = (

1881 + multimodal_embeddings[0].device

1882 + if len(multimodal_embeddings)

1883 + else mrope_positions.device

1884 + )

1885 +

1886 + # Tensors

1887 + input_ids_t = async_tensor_h2d(input_ids, device=device, dtype=torch.long)

1888 +

1889 + mm_embeddings_out = []

1890 + mm_embeddings_pos = []

1891 + # Strip position information from embeddings (last 5 channels)

1892 + # For Cohere VL, handle potentially empty frames (from unpacking)

1893 + for mm in multimodal_embeddings:

1894 + if mm.shape[0] > 0: # Only process non-empty frames

1895 + mm_embeddings_out.append(mm[:, :-5])

1896 + mm_embeddings_pos.append(mm[:, -5:].permute(1, 0).long())

1897 + else:

1898 + # Empty frame - keep as is

1899 + mm_embeddings_out.append(mm)

1900 + # Create empty position tensor with correct shape

1901 + mm_embeddings_pos.append(

1902 + torch.empty(5, 0, device=device, dtype=torch.long)

1903 + )

1904 +

1905 + positions, mrope_positions_delta = recompute_mrope_positions(

1906 + input_ids_t,

1907 + mm_embeddings_pos,

1908 + mrope_positions,

1909 + num_computed_tokens,

1910 + vision_start_token_id,

1911 + image_token_id,

1912 + video_token_id,

1913 + )

1914 +

1915 + return tuple(mm_embeddings_out), positions, mrope_positions_delta

1916 +

1917 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:

1918 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

1919 + if not mm_input_by_modality:

1920 + return None

1921 +

1922 + # The result multimodal_embeddings is tuple of tensors, with each

1923 + # tensor corresponding to a multimodal data item (image).

1924 + multimodal_embeddings: list[torch.Tensor] = []

1925 +

1926 + # NOTE: It is important to iterate over the keys in this dictionary

1927 + # to preserve the order of the modalities.

1928 + for modality in mm_input_by_modality:

1929 + multimodal_input = mm_input_by_modality[modality]

1930 + if modality == "image":

1931 + image_embeddings = self._process_image_input(multimodal_input)

1932 + image_embeddings = self._postprocess_image_embeds_evs(

1933 + image_embeddings, multimodal_input

1934 + )

1935 + multimodal_embeddings.extend(image_embeddings)

1936 +

1937 + embeddings_tuple = tuple(multimodal_embeddings)

1938 + return embeddings_tuple

1939 +

1940 + def _compute_deepstack_embeds(

1941 + self,

1942 + inputs_embeds: torch.Tensor,

1943 + multimodal_embeddings: MultiModalEmbeddings,

1944 + is_multimodal: torch.Tensor,

1945 + ) -> tuple[torch.Tensor, MultiModalEmbeddings]:

1946 + visual_lens = [len(x) for x in multimodal_embeddings]

1947 + multimodal_embeddings_cat = torch.cat(multimodal_embeddings, dim=0)

1948 +

1949 + (

1950 + multimodal_embeddings_main,

1951 + multimodal_embeddings_multiscale,

1952 + ) = torch.split(

1953 + multimodal_embeddings_cat,

1954 + [self.visual_dim, self.multiscale_dim],

1955 + dim=-1,

1956 + )

1957 +

1958 + multimodal_embeddings = torch.split(

1959 + multimodal_embeddings_main, visual_lens, dim=0

1960 + )

1961 + multimodal_embeddings_multiscale = torch.split(

1962 + multimodal_embeddings_multiscale, visual_lens, dim=0

1963 + )

1964 +

1965 + deepstack_input_embeds = inputs_embeds.new_zeros(

1966 + inputs_embeds.size(0), self.deepstack_num_level * inputs_embeds.size(1)

1967 + )

1968 +

1969 + deepstack_input_embeds = _merge_multimodal_embeddings(

1970 + inputs_embeds=deepstack_input_embeds,

1971 + multimodal_embeddings=multimodal_embeddings_multiscale,

1972 + is_multimodal=is_multimodal,

1973 + )

1974 + deepstack_input_embeds = deepstack_input_embeds.view(

1975 + inputs_embeds.shape[0], self.deepstack_num_level, self.visual_dim

1976 + )

1977 + deepstack_input_embeds = deepstack_input_embeds.permute(1, 0, 2)

1978 +

1979 + return deepstack_input_embeds, multimodal_embeddings

1980 +

1981 + def embed_input_ids(

1982 + self,

1983 + input_ids: torch.Tensor,

1984 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1985 + *,

1986 + is_multimodal: torch.Tensor | None = None,

1987 + ) -> torch.Tensor:

1988 + inputs_embeds = self._embed_text_input_ids(

1989 + input_ids,

1990 + self.language_model.embed_input_ids,

1991 + is_multimodal=is_multimodal,

1992 + )

1993 +

1994 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

1995 + return inputs_embeds

1996 +

1997 + is_multimodal = _require_is_multimodal(is_multimodal)

1998 +

1999 + if self.use_deepstack:

2000 + (

2001 + deepstack_input_embeds,

2002 + multimodal_embeddings,

2003 + ) = self._compute_deepstack_embeds(

2004 + inputs_embeds=inputs_embeds,

2005 + multimodal_embeddings=multimodal_embeddings,

2006 + is_multimodal=is_multimodal,

2007 + )

2008 + else:

2009 + deepstack_input_embeds = None

2010 +

2011 + inputs_embeds = _merge_multimodal_embeddings(

2012 + inputs_embeds=inputs_embeds,

2013 + multimodal_embeddings=multimodal_embeddings,

2014 + is_multimodal=is_multimodal,

2015 + )

2016 +

2017 + if deepstack_input_embeds is not None:

2018 + self._set_deepstack_input_embeds(deepstack_input_embeds)

2019 +

2020 + return inputs_embeds

2021 +

2022 + def forward(

2023 + self,

2024 + input_ids: torch.Tensor | None,

2025 + positions: torch.Tensor,

2026 + intermediate_tensors: IntermediateTensors | None = None,

2027 + inputs_embeds: torch.Tensor | None = None,

2028 + **kwargs: object,

2029 + ) -> torch.Tensor | IntermediateTensors:

2030 + """Run forward pass for Cohere Compass.

2031 +

2032 + Args:

2033 + input_ids: Flattened (concatenated) input_ids corresponding to a

2034 + batch.

2035 + positions: Flattened (concatenated) position ids corresponding to a

2036 + batch.

2037 + **NOTE**: If mrope is enabled (default setting for Cohere Compass

2038 + opensource models), the shape will be `(3, seq_len)`,

2039 + otherwise it will be `(seq_len,).

2040 + intermediate_tensors: Intermediate tensors from previous pipeline

2041 + stages.

2042 + inputs_embeds: Pre-computed input embeddings.

2043 + **kwargs: Additional keyword arguments including:

2044 + - pixel_values: Pixel values to be fed to a model.

2045 + `None` if no images are passed.

2046 + - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in

2047 + LLM. `None` if no images are passed.

2048 + """

2049 +

2050 + if intermediate_tensors is not None:

2051 + inputs_embeds = None

2052 +

2053 + if inputs_embeds is not None and get_pp_group().is_first_rank:

2054 + deepstack_input_embeds = self._get_deepstack_input_embeds(

2055 + inputs_embeds.size(0)

2056 + )

2057 + else:

2058 + deepstack_input_embeds = None

2059 +

2060 + hidden_states = self.language_model.model(

2061 + input_ids=input_ids,

2062 + positions=positions,

2063 + intermediate_tensors=intermediate_tensors,

2064 + inputs_embeds=inputs_embeds,

2065 + # args for deepstack

2066 + deepstack_input_embeds=deepstack_input_embeds,

2067 + )

2068 +

2069 + if inputs_embeds is not None and get_pp_group().is_first_rank:

2070 + self._clear_deepstack_input_embeds(inputs_embeds.size(0))

2071 +

2072 + return hidden_states

2073 +

2074 + def compute_logits(

2075 + self,

2076 + hidden_states: torch.Tensor,

2077 + ) -> torch.Tensor | None:

2078 + return self.language_model.compute_logits(hidden_states)

2079 +

2080 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

2081 + loader = AutoWeightsLoader(self)

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

2083 +

2084 + def get_mm_mapping(self) -> MultiModelKeys:

2085 + """

2086 + Get the module prefix in multimodal models

2087 + """

2088 + return MultiModelKeys.from_string_field(

2089 + language_model="language_model",

2090 + connector=["visual.merger", "visual.deepstack_merger_list"],

2091 + tower_model="visual.",

2092 + )

2093 +

2094 + def get_num_mm_encoder_tokens(

2095 + self,

2096 + num_image_tokens: int,

2097 + ) -> int:

2098 + hf_config = self.config

2099 + vision_config = hf_config.vision_config

2100 + merge_size = vision_config.spatial_merge_size

2101 +

2102 + return num_image_tokens * merge_size**2

2103 +

2104 + def get_num_mm_connector_tokens(

2105 + self,

2106 + num_vision_tokens: int,

2107 + ) -> int:

2108 + hf_config = self.config

2109 + vision_config = hf_config.vision_config

2110 + merge_size = vision_config.spatial_merge_size

2111 + return num_vision_tokens // merge_size**2

2112 +

2113 +

2114 +@lru_cache

2115 +def _cached_tensor(x, device) -> torch.Tensor:

2116 + return torch.tensor(x, device=device)

2117 +

2118 +

2119 +class CohereCompassAttention(nn.Module):

2120 + def __init__(

2121 + self,

2122 + config,

2123 + cache_config: CacheConfig | None = None,

2124 + quant_config: QuantizationConfig | None = None,

2125 + max_position_embeddings: int | None = None,

2126 + prefix: str = "",

2127 + ):

2128 + super().__init__()

2129 + tp_size = get_tensor_model_parallel_world_size()

2130 + layer_idx = extract_layer_index(prefix)

2131 + layer_type = config.layer_types[layer_idx]

2132 +

2133 + self.head_dim = config.head_dim

2134 + self.total_num_heads = config.num_attention_heads

2135 + self.num_heads = self.total_num_heads // tp_size

2136 + self.total_num_kv_heads = config.num_key_value_heads

2137 + if self.total_num_kv_heads >= tp_size:

2138 + assert self.total_num_kv_heads % tp_size == 0

2139 + else:

2140 + assert tp_size % self.total_num_kv_heads == 0

2141 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

2142 + self.q_size = self.num_heads * self.head_dim

2143 + self.kv_size = self.num_kv_heads * self.head_dim

2144 + self.scaling = self.head_dim**-0.5

2145 +

2146 + self.qkv_proj = QKVParallelLinear(

2147 + config.hidden_size,

2148 + self.head_dim,

2149 + self.total_num_heads,

2150 + self.total_num_kv_heads,

2151 + bias=config.attention_bias,

2152 + quant_config=quant_config,

2153 + prefix=f"{prefix}.qkv_proj",

2154 + )

2155 + self.o_proj = RowParallelLinear(

2156 + self.total_num_heads * self.head_dim,

2157 + config.hidden_size,

2158 + bias=config.attention_bias,

2159 + quant_config=quant_config,

2160 + prefix=f"{prefix}.o_proj",

2161 + )

2162 +

2163 + rope_parameters = config.rope_parameters.get(layer_type)

2164 + self.rotary_emb = (

2165 + get_rope(

2166 + self.head_dim,

2167 + max_position=max_position_embeddings or config.max_position_embeddings,

2168 + rope_parameters=rope_parameters,

2169 + is_neox_style=True,

2170 + )

2171 + if rope_parameters is not None

2172 + else None

2173 + )

2174 + sliding_window = (

2175 + config.sliding_window if layer_type == "sliding_attention" else None

2176 + )

2177 + self.attn = Attention(

2178 + self.num_heads,

2179 + self.head_dim,

2180 + self.scaling,

2181 + num_kv_heads=self.num_kv_heads,

2182 + cache_config=cache_config,

2183 + quant_config=quant_config,

2184 + per_layer_sliding_window=sliding_window,

2185 + prefix=f"{prefix}.attn",

2186 + )

2187 +

2188 + def forward(

2189 + self,

2190 + positions: torch.Tensor,

2191 + hidden_states: torch.Tensor,

2192 + ) -> torch.Tensor:

2193 + qkv, _ = self.qkv_proj(hidden_states)

2194 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

2195 + if self.rotary_emb is not None:

2196 + q, k = self.rotary_emb(positions, q, k)

2197 + attn_output = self.attn(q, k, v)

2198 + output, _ = self.o_proj(attn_output)

2199 + return output

2200 +

2201 +

2202 +class CohereCompassDecoderLayer(nn.Module):

2203 + def __init__(

2204 + self,

2205 + config,

2206 + cache_config: CacheConfig | None = None,

2207 + quant_config: QuantizationConfig | None = None,

2208 + max_position_embeddings: int | None = None,

2209 + prefix: str = "",

2210 + ):

2211 + super().__init__()

2212 + self.self_attn = CohereCompassAttention(

2213 + config,

2214 + cache_config,

2215 + quant_config=quant_config,

2216 + max_position_embeddings=max_position_embeddings,

2217 + prefix=f"{prefix}.self_attn",

2218 + )

2219 + self.mlp = CohereMLP(

2220 + config,

2221 + quant_config=quant_config,

2222 + reduce_results=True,

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

2224 + )

2225 + self.input_layernorm = LayerNorm(

2226 + param_shape=config.hidden_size,

2227 + eps=config.layer_norm_eps,

2228 + )

2229 +

2230 + def forward(

2231 + self,

2232 + positions: torch.Tensor,

2233 + hidden_states: torch.Tensor,

2234 + residual: torch.Tensor | None,

2235 + ) -> tuple[torch.Tensor, torch.Tensor]:

2236 + residual = hidden_states

2237 + hidden_states, residual = self.input_layernorm(hidden_states, residual)

2238 + hidden_states_attention = self.self_attn(

2239 + positions=positions,

2240 + hidden_states=hidden_states,

2241 + )

2242 + hidden_states_mlp = self.mlp(hidden_states)

2243 + hidden_states = residual + hidden_states_attention + hidden_states_mlp

2244 + return hidden_states, residual

2245 +

2246 +

2247 +@support_torch_compile(

2248 + dynamic_arg_dims={

2249 + "input_ids": 0,

2250 + "positions": -1,

2251 + "intermediate_tensors": 0,

2252 + "inputs_embeds": 0,

2253 + "deepstack_input_embeds": 0,

2254 + }

2255 +)

2256 +class CohereCompassTextModel(nn.Module):

2257 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

2258 + super().__init__()

2259 + config = vllm_config.model_config.hf_config

2260 + cache_config = vllm_config.cache_config

2261 + quant_config = vllm_config.quant_config

2262 +

2263 + self.config = config

2264 + self.vocab_size = config.vocab_size

2265 + self.embed_tokens = VocabParallelEmbedding(

2266 + config.vocab_size,

2267 + config.hidden_size,

2268 + )

2269 + self.start_layer, self.end_layer, self.layers = make_layers(

2270 + config.num_hidden_layers,

2271 + lambda prefix: CohereCompassDecoderLayer(

2272 + config,

2273 + cache_config,

2274 + quant_config=quant_config,

2275 + max_position_embeddings=vllm_config.model_config.max_model_len,

2276 + prefix=prefix,

2277 + ),

2278 + prefix=f"{prefix}.layers",

2279 + )

2280 + self.norm = LayerNorm(

2281 + param_shape=config.hidden_size,

2282 + eps=config.layer_norm_eps,

2283 + )

2284 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

2285 + ["hidden_states", "residual"], config.hidden_size

2286 + )

2287 +

2288 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

2289 + return self.embed_tokens(input_ids)

2290 +

2291 + def forward(

2292 + self,

2293 + input_ids: torch.Tensor | None,

2294 + positions: torch.Tensor,

2295 + intermediate_tensors: IntermediateTensors | None = None,

2296 + inputs_embeds: torch.Tensor | None = None,

2297 + deepstack_input_embeds: IntermediateTensors | None = None,

2298 + ) -> torch.Tensor | IntermediateTensors:

2299 + if get_pp_group().is_first_rank:

2300 + hidden_states = (

2301 + inputs_embeds

2302 + if inputs_embeds is not None

2303 + else self.embed_input_ids(input_ids)

2304 + )

2305 + residual = None

2306 + else:

2307 + assert intermediate_tensors is not None

2308 + hidden_states = intermediate_tensors["hidden_states"]

2309 + residual = intermediate_tensors["residual"]

2310 +

2311 + for layer_idx, layer in islice(

2312 + enumerate(self.layers), self.start_layer, self.end_layer

2313 + ):

2314 + hidden_states, residual = layer(

2315 + positions,

2316 + hidden_states,

2317 + residual,

2318 + )

2319 + if deepstack_input_embeds is not None and layer_idx < len(

2320 + deepstack_input_embeds

2321 + ):

2322 + hidden_states = (

2323 + hidden_states

2324 + + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]

2325 + )

2326 +

2327 + if not get_pp_group().is_last_rank:

2328 + return IntermediateTensors(

2329 + {"hidden_states": hidden_states, "residual": residual}

2330 + )

2331 + hidden_states, _ = self.norm(hidden_states, residual)

2332 + return hidden_states

2333 +

2334 +

2335 +class CohereCompassForCausalLM(CohereForCausalLM):

2336 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

2337 + nn.Module.__init__(self)

2338 + config = vllm_config.model_config.hf_config

2339 + assert config.tie_word_embeddings

2340 +

2341 + self.config = config

2342 + self.quant_config = vllm_config.quant_config

2343 + logit_scale = 1.0 if config.logit_scale is None else config.logit_scale

2344 + self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)

2345 + self.model = CohereCompassTextModel(

2346 + vllm_config=vllm_config,

2347 + prefix=maybe_prefix(prefix, "model"),

2348 + )

2349 + self.make_empty_intermediate_tensors = (

2350 + self.model.make_empty_intermediate_tensors

2351 + )

2352 +

2353 +

2354 +__all__ = ["CohereCompassForConditionalGeneration"]

[Model] Add Cohere Compass model (#54774)