update test converage to FA2 and FA3

Signed-off-by: qqma <qqma@amazon.com>

Author: qqma

tests/compile/piecewise/test_full_cudagraph.py

@@ -46,7 +46,10 @@ class BackendConfig:
     # FA3 on Hopper
     "FA3":
     BackendConfig(name="FA3",
-                  env_vars={"VLLM_FLASH_ATTN_VERSION": "3"},
+                  env_vars={
+                      "VLLM_FLASH_ATTN_VERSION": "3",
+                      "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": 16,
+                  },
                   comp_config={
                       "cudagraph_mode": "FULL",
                   },
@@ -90,7 +93,10 @@ class BackendConfig:
     # FA2
     "FA2":
     BackendConfig(name="FA2",
-                  env_vars={"VLLM_FLASH_ATTN_VERSION": "2"},
+                  env_vars={
+                      "VLLM_FLASH_ATTN_VERSION": "2",
+                      "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": 16,
+                  },
                   comp_config={
                       "cudagraph_mode": "FULL",
                   }),

tests/v1/cudagraph/test_cudagraph_mode.py

@@ -47,7 +47,10 @@ class BackendConfig:
     # FA3 on Hopper
     "FA3":
     BackendConfig(name="FA3",
-                  env_vars={"VLLM_FLASH_ATTN_VERSION": "3"},
+                  env_vars={
+                      "VLLM_FLASH_ATTN_VERSION": "3",
+                      "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": 16,
+                  },
                   comp_config={
                       "cudagraph_mode": "FULL",
                   },
@@ -76,7 +79,10 @@ class BackendConfig:
     # FA2
     "FA2":
     BackendConfig(name="FA2",
-                  env_vars={"VLLM_FLASH_ATTN_VERSION": "2"},
+                  env_vars={
+                      "VLLM_FLASH_ATTN_VERSION": "2",
+                      "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": 16,
+                  },
                   comp_config={
                       "cudagraph_mode": "FULL_AND_PIECEWISE",
                   }),