import torch
from transformers.models.llava.modeling_llava import LlavaForConditionalGeneration


class TestLlavaForConditionalGeneration(LlavaForConditionalGeneration):
    def forward(
        self,
        input_ids= None,
        pixel_values= None,
        attention_mask = None,
        position_ids= None,
        past_key_values= None,
        inputs_embeds= None,
        vision_feature_layer= None,
        vision_feature_select_strategy= None,
        labels= None,
        use_cache= None,
        output_attentions= None,
        output_hidden_states= None,
        return_dict= None,
    ):
        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
        vision_feature_layer = (
            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
        )
        vision_feature_select_strategy = (
            vision_feature_select_strategy
            if vision_feature_select_strategy is not None
            else self.config.vision_feature_select_strategy
        )

        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 input_ids.shape[1] != 1:
                image_outputs = self.vision_tower(pixel_values, output_hidden_states=True)
                # this is not memory efficient at all (output_hidden_states=True) will save all the hidden stated.
                selected_image_feature = image_outputs.hidden_states[vision_feature_layer]

                if vision_feature_select_strategy == "default":
                    selected_image_feature = selected_image_feature[:, 1:]
                elif vision_feature_select_strategy == "full":
                    selected_image_feature = selected_image_feature
                else:
                    raise ValueError(
                        f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}"
                    )

                image_features = self.multi_modal_projector(selected_image_feature)
                inputs_embeds, attention_mask, position_ids = self._merge_input_ids_with_image_features(
                    image_features, inputs_embeds, input_ids, attention_mask, position_ids
                )
                if labels is None:
                    labels = torch.full_like(attention_mask, self.config.ignore_index).to(torch.long)
            else:
                # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
                # generation with cache
                if 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, :, 0]
                    batch_index, non_attended_tokens = torch.where(first_layer_past_key_value == 0)

                    ############################
                    # Add here
                    # non_attended_tokens = non_attended_tokens - attention_mask.shape[1]
                    ############################

                    # Get the target length
                    target_seqlen = first_layer_past_key_value.shape[-1] + 1

                    extended_attention_mask = torch.ones(
                        (attention_mask.shape[0], target_seqlen - attention_mask.shape[1]),
                        dtype=attention_mask.dtype,
                        device=attention_mask.device,
                    )

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

                    attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1)
                    position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
        print(position_ids)

llava = TestLlavaForConditionalGeneration.from_pretrained(
    "llava-hf/llava-1.5-13b-hf", torch_dtype=torch.float16)
llava.to("cuda")
input_ids = torch.tensor([[319]]).cuda().long()
pixel_values = torch.zeros((1, 3, 224, 224)).cuda().half()
past_key_values = torch.rand((1, 1, 1, 40, 642, 128)).cuda().half()
past_key_values[:, :, :, 0, 600, 0] = 0.
attention_mask = torch.ones((1, 68)).cuda().half()
llava(input_ids=input_ids,
      pixel_values=pixel_values,
      past_key_values=past_key_values,
      attention_mask=attention_mask)

Traceback (most recent call last):
  File "/workspace/a.py", line 99, in <module>
    llava(input_ids=input_ids,
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/workspace/a.py", line 87, in forward
    attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1)
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

batch_index, non_attended_tokens = torch.where(first_layer_past_key_value == 0)
# Get the target length
target_seqlen = first_layer_past_key_value.shape[-1] + 1

# Now the index of `non_attended_tokens` corresponds to the index of `target_seqlen` == first_layer_past_key_value.axis(-1).

# However, the index of extended_attention_mask is target_seqlen - attention_mask.shape[1].
# Are the indices of `non_attended_tokens` and `extended_attention_mask` different?
extended_attention_mask = torch.ones(
    (attention_mask.shape[0], target_seqlen - attention_mask.shape[1]),
    dtype=attention_mask.dtype,
    device=attention_mask.device,
)

# Zero-out the places where we don't need to attend
extended_attention_mask[batch_index, non_attended_tokens] = 0
attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1)