add VLLM_USE_PA_PRINT_PARAM to print v1 fa size

vllm/v1/attention/backends/flash_attn.py

@@ -7,6 +7,7 @@ from typing import ClassVar, Optional
 import numpy as np
 import torch
 
+import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionMetadata, AttentionType,
@@ -521,6 +522,11 @@ class FlashAttentionImpl(AttentionImpl):
                     num_splits=attn_metadata.max_num_splits,
                 )
             else:
+                if envs.VLLM_USE_PA_PRINT_PARAM:
+                    print("PA SIZE:")
+                    print(f"q.shape = {query[:num_actual_tokens].unsqueeze(1).shape}, key_cache.shape = {key_cache.shape}, value_cache.shape = {value_cache.shape}")
+                    print(f"cu_seqlens_q.shape = {cu_seqlens_q.shape}, max_seqlen_q = {max_seqlen_q}, seqused_k.shape = {seqused_k.shape}, max_seqlen_k = {max_seqlen_k}")
+                    print(f"softmax_scale = {self.scale:.3f}, alibi_slopes = {self.alibi_slopes}, window_size = {self.sliding_window}, block_tables.shape = {block_table.shape}, softcap = {self.logits_soft_cap}, scheduler_metadata = {scheduler_metadata}")
                 vllm_flash_attn_varlen_func(
                     q=query[:num_actual_tokens],
                     k=key_cache,