vllm.models.glm5next.nvidia.multimodal ¶

class Glm5NextMultiModalProcessor(Glm4vMultiModalProcessor):
    """The vLLM-native ``Glm5NextProcessor`` extracts image/video features
    only and passes the prompt text through unchanged, so prompt expansion
    (image token repeat, video frame/timestamp structure) is owned by vLLM's
    prompt-update machinery — the inherited ``_get_prompt_updates`` builds
    the replacement content and the placeholder scan validates against
    exactly that."""

    def _hf_processor_applies_updates(
        self,
        prompt_text: str,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        tokenization_kwargs: Mapping[str, object],
    ) -> bool:
        return False

class Glm5NextProcessingInfo(Glm4vProcessingInfo):
    """Wires up the vLLM-native processor for the multimodal checkpoint.

    The checkpoint's ``processor_config.json`` declares a custom ``processor_class``
    and stores its image/video processor configs inline (no standalone
    ``preprocessor_config.json``), so ``AutoProcessor`` cannot resolve the
    config. We bypass it and build our own ``Glm5NextProcessor``
    (``vllm/transformers_utils/processors/glm5next.py``), a port of the
    training-side pipeline that no longer imports transformers' GLM processor
    classes. The port applies ``patch_expand_factor`` (checkpoint ships 1)
    inside ``smart_resize``'s spatial factor.
    """

    @cached_property
    def _glm5_hf_processor(self):
        from vllm.transformers_utils.processors.glm5next import Glm5NextProcessor

        return Glm5NextProcessor.from_pretrained(self.ctx.model_config.model)

    def get_hf_processor(self, **kwargs: object):
        return self._glm5_hf_processor

    def _processor_pixel_budget(self, proc) -> tuple[int, int]:
        from vllm.transformers_utils.processors.glm5next import _pixel_budget

        return _pixel_budget(
            proc.min_image_tokens,
            proc.max_image_tokens,
            proc.patch_size,
            proc.merge_size,
            proc.temporal_patch_size,
        )

    def _get_image_max_pixels(self) -> int:
        mm_kwargs = self.ctx.get_merged_mm_kwargs({})
        if (override := mm_kwargs.get("max_pixels")) is not None:
            return int(override)
        return self._processor_pixel_budget(self.get_hf_processor().image_processor)[1]

    def _get_video_max_pixels(self) -> int:
        mm_kwargs = self.ctx.get_merged_mm_kwargs({})
        if (override := mm_kwargs.get("max_pixels")) is not None:
            return int(override)
        return self._processor_pixel_budget(self.get_hf_processor().video_processor)[1]

    def _get_vision_info(
        self,
        *,
        image_width: int,
        image_height: int,
        num_frames: int = 16,
        do_resize: bool = True,
        max_image_pixels: int = 28 * 28 * 2 * 30000,
    ) -> tuple[ImageSize, int]:
        """GLM-5.3-Flash canvas geometry for token budgeting and dummy inputs.

        The inherited Glm4v path resolves the pixel budget from
        ``size.longest_edge`` and resizes with GLM-4V's ``smart_resize``. This
        checkpoint's ``processor_config.json`` ships the token-budget style
        (``min_image_tokens`` / ``max_image_tokens``) with no ``size`` key, and
        the alignment factor carries ``patch_expand_factor`` — resolve both
        from the vLLM-native processor so profiling matches runtime geometry.
        """
        from vllm.transformers_utils.processors.glm5next import smart_resize

        vision_config = self.get_hf_config().vision_config
        patch_size = vision_config.patch_size
        merge_size = vision_config.spatial_merge_size
        temporal_patch_size = vision_config.temporal_patch_size

        image_processor = self.get_hf_processor().image_processor
        factor = patch_size * merge_size * image_processor.patch_expand_factor
        # Keep the profiling search viable when the caller's budget is below
        # one aligned canvas of the requested duration.
        max_image_pixels = max(max_image_pixels, temporal_patch_size * factor * factor)

        if do_resize:
            t = num_frames if num_frames > temporal_patch_size else temporal_patch_size
            resized_height, resized_width = smart_resize(
                t=t,
                h=image_height,
                w=image_width,
                t_factor=temporal_patch_size,
                h_factor=factor,
                w_factor=factor,
                min_pixels=1,
                max_pixels=max_image_pixels,
            )
            preprocessed_size = ImageSize(width=resized_width, height=resized_height)
        else:
            preprocessed_size = ImageSize(width=image_width, height=image_height)

        padded_num_frames = num_frames + (-num_frames % temporal_patch_size)
        grid_t = max(padded_num_frames // temporal_patch_size, 1)
        grid_h = preprocessed_size.height // patch_size
        grid_w = preprocessed_size.width // patch_size

        num_patches = grid_t * grid_h * grid_w
        num_vision_tokens = num_patches // (merge_size**2)

        return preprocessed_size, num_vision_tokens

class Glm5NextVisionTransformer(nn.Module):
    # Stacked-weight remap for the GLM-OCR/GLM-4V vision checkpoint layout.
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_stacked={
            ".attn.q.": (".attn.qkv.", "q"),
            ".attn.k.": (".attn.qkv.", "k"),
            ".attn.v.": (".attn.qkv.", "v"),
            ".gate_proj": (".gate_up_proj", 0),
            ".up_proj": (".gate_up_proj", 1),
        }
    )

    def __init__(
        self,
        text_config,  # noqa: ANN001
        vision_config,
        norm_eps: float = 1e-6,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        use_data_parallel = is_vit_use_data_parallel()
        self.tp_size = (
            1 if use_data_parallel else get_tensor_model_parallel_world_size()
        )

        patch_size = vision_config.patch_size
        temporal_patch_size = vision_config.temporal_patch_size
        in_channels = vision_config.in_channels
        depth = vision_config.depth
        self.hidden_size = vision_config.hidden_size
        self.num_heads = vision_config.num_heads

        self.patch_size = vision_config.patch_size
        self.spatial_merge_size = vision_config.spatial_merge_size
        self.out_hidden_size = vision_config.out_hidden_size

        swiglu_limit = vision_config.swiglu_limit
        if swiglu_limit is None:
            swiglu_limit = text_config.swiglu_limit
        assert swiglu_limit is not None, (
            "GLM-5.3-Flash vision requires swiglu_limit (vision_config or text_config)"
        )

        # Single construction pass — no abs-pos embeddings / post-conv norm (OCR delta).
        self.patch_embed = Glm5NextVisionPatchEmbed(
            patch_size=patch_size,
            temporal_patch_size=temporal_patch_size,
            in_channels=in_channels,
            hidden_size=self.hidden_size,
        )

        norm_layer = partial(RMSNorm, eps=norm_eps)
        head_dim = self.hidden_size // self.num_heads
        self.rotary_pos_emb = get_rope(
            head_size=head_dim,
            max_position=8192,
            is_neox_style=True,
            rope_parameters={"partial_rotary_factor": 0.5},
        )
        self.blocks = nn.ModuleList(
            [
                Glm5NextVisionBlock(
                    dim=self.hidden_size,
                    num_heads=self.num_heads,
                    mlp_hidden_dim=vision_config.intermediate_size,
                    swiglu_limit=swiglu_limit,
                    norm_layer=norm_layer,
                    quant_config=quant_config,
                    prefix=f"{prefix}.blocks.{layer_idx}",
                )
                for layer_idx in range(depth)
            ]
        )
        # GLM-5.3-Flash merger bottleneck width.
        self.merger = Glm5NextPatchMerger(
            d_model=vision_config.out_hidden_size,
            context_dim=vision_config.projection_intermediate_size,
            swiglu_limit=swiglu_limit,
            quant_config=quant_config,
            bias=False,
            prefix=f"{prefix}.merger",
        )

        self.downsample = Conv2dLayer(
            in_channels=vision_config.hidden_size,
            out_channels=vision_config.out_hidden_size,
            kernel_size=vision_config.spatial_merge_size,
            stride=vision_config.spatial_merge_size,
        )
        self.post_layernorm = RMSNorm(
            vision_config.hidden_size, eps=vision_config.rms_norm_eps
        )

        self.attn_backend = get_vit_attn_backend(
            head_size=head_dim,
            dtype=torch.get_default_dtype(),
        )

    @property
    def dtype(self) -> torch.dtype:
        return self.patch_embed.proj.weight.dtype

    @property
    def device(self) -> torch.device:
        return self.patch_embed.proj.weight.device

    def rot_pos_emb(
        self, grid_thw: list[list[int]]
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        pos_ids = []
        for t, h, w in grid_thw:
            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
            hpos_ids = (
                hpos_ids.reshape(
                    h // self.spatial_merge_size,
                    self.spatial_merge_size,
                    w // self.spatial_merge_size,
                    self.spatial_merge_size,
                )
                .permute(0, 2, 1, 3)
                .flatten()
            )
            wpos_ids = (
                wpos_ids.reshape(
                    h // self.spatial_merge_size,
                    self.spatial_merge_size,
                    w // self.spatial_merge_size,
                    self.spatial_merge_size,
                )
                .permute(0, 2, 1, 3)
                .flatten()
            )
            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
        pos_ids = torch.cat(pos_ids, dim=0)
        max_grid_size = max(max(h, w) for _, h, w in grid_thw)

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

        pos_ids = pos_ids.to(cos.device, non_blocking=True)
        cos_combined = cos[pos_ids].flatten(1)
        sin_combined = sin[pos_ids].flatten(1)
        return cos_combined, sin_combined, pos_ids

    def compute_attn_mask_seqlen(
        self,
        cu_seqlens: torch.Tensor,
    ) -> torch.Tensor | None:
        max_seqlen = None
        if self.attn_backend in {
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.ROCM_AITER_FA,
            AttentionBackendEnum.TRITON_ATTN,
        }:
            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
        return max_seqlen

    def prepare_encoder_metadata(
        self,
        grid_thw_list: list[list[int]],
        *,
        max_batch_size: int | None = None,
        max_frames_per_batch: int | None = None,
        max_seqlen_override: int | None = None,
        device: torch.device | None = None,
    ) -> dict[str, torch.Tensor | None]:
        """Compute encoder metadata for eager and CUDA graph execution."""
        if device is None:
            device = self.device

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

        rotary_cos, rotary_sin, _ = self.rot_pos_emb(grid_thw_list)
        metadata["rotary_pos_emb_cos"] = rotary_cos
        metadata["rotary_pos_emb_sin"] = rotary_sin

        grid_thw_np = np.array(grid_thw_list, dtype=np.int32)
        patches_per_frame = grid_thw_np[:, 1] * grid_thw_np[:, 2]
        cu_seqlens = np.repeat(patches_per_frame, grid_thw_np[:, 0]).cumsum(
            dtype=np.int32
        )
        cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])

        pad_to = (
            max_frames_per_batch if max_frames_per_batch is not None else max_batch_size
        )
        if pad_to is not None:
            num_seqs = len(cu_seqlens) - 1
            if num_seqs < pad_to:
                cu_seqlens = np.concatenate(
                    [
                        cu_seqlens,
                        np.full(
                            pad_to - num_seqs,
                            cu_seqlens[-1],
                            dtype=np.int32,
                        ),
                    ]
                )

        metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(
            self.attn_backend, cu_seqlens, device
        )

        if max_seqlen_override is not None:
            max_seqlen_val = max_seqlen_override
        else:
            max_seqlen_val = MMEncoderAttention.compute_max_seqlen(
                self.attn_backend, cu_seqlens
            )
        metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)

        metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
            self.attn_backend,
            cu_seqlens,
            self.hidden_size,
            self.tp_size,
            device,
        )

        return metadata

    def forward(
        self,
        x: torch.Tensor,
        grid_thw: torch.Tensor | list[list[int]],
        *,
        encoder_metadata: dict[str, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        # patchify
        x = x.to(device=self.device, dtype=self.dtype)
        x = self.patch_embed(x)

        if encoder_metadata is not None:
            # Encoder CUDA-graph path (PR #49852): rotary/cu_seqlens/max_seqlen are
            # precomputed by prepare_encoder_metadata (which uses rot_pos_emb exactly
            # as the eager rebuild does), so reuse them and skip the per-call CPU
            # rebuild (the low-GPU-util culprit on multimodal workloads).
            rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"]
            rotary_pos_emb_sin = encoder_metadata["rotary_pos_emb_sin"]
            cu_seqlens = encoder_metadata["cu_seqlens"]
            max_seqlen = encoder_metadata["max_seqlen"]
        else:
            if isinstance(grid_thw, list):
                grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
            rotary_pos_emb_cos, rotary_pos_emb_sin, _ = self.rot_pos_emb(grid_thw)
            cu_seqlens = torch.repeat_interleave(
                grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
            ).cumsum(dim=0, dtype=torch.int32)
            cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
            cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
            max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

        # transformers
        x = x.unsqueeze(1)
        for blk in self.blocks:
            x = blk(
                x,
                cu_seqlens=cu_seqlens,
                rotary_pos_emb_cos=rotary_pos_emb_cos,
                rotary_pos_emb_sin=rotary_pos_emb_sin,
                max_seqlen=max_seqlen,
            )

        # adapter
        x = self.post_layernorm(x)
        x = x.view(-1, self.spatial_merge_size, self.spatial_merge_size, x.shape[-1])
        x = x.permute(0, 3, 1, 2)
        x = self.downsample(x).view(-1, self.out_hidden_size)
        x = self.merger(x)
        return x

    def load_weights(self, weights) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)