llmcompressor.modeling.kimi_k3.modeling_kimi_k3

def forward(
    self,
    input_ids: torch.LongTensor | None = None,
    pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None,
    grid_thws: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: list[torch.FloatTensor] | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
) -> tuple | LlavaCausalLMOutputWithPast:
    r"""
    Args:
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

    ```"""
    assert self.vision_tower is not None, "vision_tower is not loaded"
    output_attentions = (
        output_attentions
        if output_attentions is not None
        else self.config.output_attentions
    )
    output_hidden_states = (
        output_hidden_states
        if output_hidden_states is not None
        else self.config.output_hidden_states
    )
    return_dict = (
        return_dict if return_dict is not None else self.config.use_return_dict
    )

    if inputs_embeds is None:
        # 1. Extra the input embeddings
        inputs_embeds = self.get_input_embeddings()(input_ids)

        # 2. Merge text and images
        if (
            pixel_values is not None
            and len(pixel_values) > 0
            and input_ids.shape[1] != 1
        ):
            image_features = self._extract_image_features(pixel_values, grid_thws)
            if self.mm_projector:
                image_features = self.mm_projector(image_features)

            inputs_embeds = inputs_embeds.to(
                image_features[0].dtype
            )  # num_tokens, embed_dim
            inputs_embeds, attention_mask, labels, position_ids = (
                self._merge_input_ids_with_image_features(
                    image_features,
                    inputs_embeds,
                    input_ids,
                    attention_mask,
                    labels,
                )
            )

        # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
        # generation with cache
        elif (
            past_key_values is not None
            and pixel_values is not None
            and input_ids.shape[1] == 1
        ):
            # Retrieve the first layer to inspect the logits and mask out the hidden states
            # that are set to 0
            first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]

            # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
            batch_index, non_attended_tokens = torch.where(
                first_layer_past_key_value.float().sum(-2) == 0
            )

            # Get the target length
            target_length = input_ids.shape[1]
            past_length = first_layer_past_key_value.shape[-1]

            extended_attention_mask = torch.ones(
                (attention_mask.shape[0], past_length),
                dtype=attention_mask.dtype,
                device=attention_mask.device,
            )

            # Filter out only the tokens that can be un-attended, this can happen
            # if one uses Llava + Fused modules where the cache on the
            # first iteration is already big enough, or if one passes custom cache
            valid_indices = non_attended_tokens < extended_attention_mask.size(-1)
            new_batch_index = batch_index[valid_indices]
            new_non_attended_tokens = non_attended_tokens[valid_indices]

            # Zero-out the places where we don't need to attend
            extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0

            attention_mask = torch.cat(
                (extended_attention_mask, attention_mask[:, -target_length:]), dim=1
            )
            position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1

    outputs = self.language_model(
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    logits = outputs[0]

    loss = None
    if labels is not None:
        # Shift so that tokens < n predict n
        if attention_mask is not None:
            shift_attention_mask = attention_mask[..., 1:]
            shift_logits = logits[..., :-1, :][
                shift_attention_mask.to(logits.device) != 0
            ].contiguous()
            shift_labels = labels[..., 1:][
                shift_attention_mask.to(labels.device) != 0
            ].contiguous()
        else:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
        # Flatten the tokens
        loss_fct = nn.CrossEntropyLoss()
        loss = loss_fct(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1).to(shift_logits.device),
        )

    if not return_dict:
        output = (logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output

    return LlavaCausalLMOutputWithPast(
        loss=loss,
        logits=logits,
        past_key_values=outputs.past_key_values,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

def forward(
    self, pixel_values: torch.Tensor, grid_thws: torch.Tensor
) -> torch.Tensor:
    """
    Args:
        pixel_values (torch.Tensor): The input pixel values.
        grid_thws (torch.Tensor): Temporal, height and width.

    Returns:
        torch.Tensor: The output tokens.
    """
    # grid_thws = grid_thws.to('cpu')
    assert grid_thws.ndim == 2, f"grid_thws should be 2D, got {grid_thws.ndim}"
    assert grid_thws.size(1) == 3, f"No support for thw: {grid_thws}"
    hidden_states = self.patch_embed(pixel_values, grid_thws)
    hidden_states = self.encoder(hidden_states, grid_thws)
    if (
        self.merge_type == "sd2_tpool"
    ):  # spatial downsampling 2x with temporal pooling all
        hidden_states = tpool_patch_merger(
            hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size
        )
    else:
        raise NotImplementedError(f"Not support {self.merge_type}")

    return hidden_states

def attention_qkvpacked(
    self,
    x: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: torch.Tensor,
    rope_freqs_cis: torch.Tensor | None = None,
):
    """
    Args:
        x (torch.Tensor): (batch_size, seqlen, hidden_dim)
        cu_seqlens (torch.Tensor):
    """
    xqkv = self.wqkv(x)

    qkv_shape = xqkv.size()[:-1] + (
        3,
        self.num_heads,
        self.hidden_size_per_attention_head,
    )
    # xqkv: (batch_size, seqlen, 3, nheads, headdim)
    xqkv = xqkv.view(*qkv_shape)
    xq, xk, xv = torch.unbind(xqkv, dim=-3)

    xq, xk = apply_rope(xq, xk, rope_freqs_cis)

    attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
    attn_out = attn_func(
        xq,
        xk,
        xv,
        q_cu_seqlens=cu_seqlens,
        k_cu_seqlens=cu_seqlens,
        max_seqlen_k=max_seqlen,
        max_seqlen_q=max_seqlen,
        deterministic=self.use_deterministic_attn,
    )

    attn_out = self.wo(attn_out)
    return attn_out

def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor:
    """
    Args:
        x (L, Channels): input tensor
        grid_hws (N, 3): temporal, height and width

    Returns:
        (L, Cout) tensor
    """
    x = self.proj(x).view(x.size(0), -1)
    # apply positional embedding
    x = self.pos_emb(x, grid_thws)
    return x

def get_freqs_cis(
    self, grid_thws: torch.Tensor, device: torch.device
) -> torch.Tensor:
    """
    Args:
        grid_thws (torch.Tensor): grid time, height and width

    Returns:
        freqs_cis: tensor of shape (sum(t * height * width), dim//2)
    """
    if not hasattr(self, "freqs_cis"):
        self.register_buffer(
            "freqs_cis", self._precompute_freqs_cis(device), persistent=False
        )

    shapes = grid_thws.tolist()
    assert all(
        1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes
    ), (
        shapes,
        self.max_height,
        self.max_width,
    )
    freqs_cis = torch.cat(
        [
            self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1)
            for t, h, w in shapes
        ],
        dim=0,
    )
    return freqs_cis