How to get seq_len in Atten

[lucasjinreal]

def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: KVCache,
        input_metadata: InputMetadata,
        cache_event: Optional[torch.cuda.Event],
    ) -> torch.Tensor:
        qkv, _ = self.W_pack(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
        k_cache, v_cache = kv_cache
# How to get seq_len here?
        attn_output = self.attn(
            positions, q, k, v, k_cache, v_cache, input_metadata, cache_event, seq_len=512
        )
        output, _ = self.o_proj(attn_output)
        return output

How to get seq_len at Atten layer. (I printed shapes of each tensor but can't really know how to according them to torch side), Can anyone help?

[zhuohan123]

You can check input_metadata.prompt_lens and input_metadata.context_lens. They are the sequence lengths for prompt and generation tokens respectively.

[lucasjinreal]

@zhuohan123 thank u. so the context_lens here should equalent to:

 kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            kv_seq_len += past_key_value[0].shape[-2]
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

kv_seq_len here in torch?