[torch.compile] add a flag to track batchsize statistics (#11059)

Signed-off-by: youkaichao <youkaichao@gmail.com>
2024-12-10 12:40:52 -08:00
parent e739194926
commit 75f89dc44c
4 changed files with 37 additions and 1 deletions
--- a/vllm/v1/attention/backends/flash_attn.py
+++ b/vllm/v1/attention/backends/flash_attn.py
@@ -56,6 +56,7 @@ class FlashAttentionMetadata:
    seq_start_loc: torch.Tensor
    block_table: torch.Tensor
    slot_mapping: torch.Tensor
+    num_input_tokens: int = 0  # Number of tokens including padding.


 class FlashAttentionImpl(AttentionImpl):