support FULL graph mode for GQA (#3970)

### What this PR does / why we need it?
The current library only supports the FullDecodeOnly graph mode, which
enables full graph execution during the decode. This PR extends support
to allow full graph execution in both the prefill and decode, referred
to as FULL graph mode.

- vLLM version: v0.11.0
- vLLM main:
vllm-project/vllm@2918c1b

Signed-off-by: wangxiaoxin-sherie <[email protected]>
Co-authored-by: wangxiaoxin-sherie <[email protected]>

Author: momo609

.github/workflows/_e2e_test.yaml

@@ -180,6 +180,7 @@ jobs:
         if: ${{ inputs.type == 'full' }}
         run: |
           pytest -sv tests/e2e/multicard/test_torchair_graph_mode.py
+          pytest -sv tests/e2e/multicard/test_full_graph_mode.py
           pytest -sv tests/e2e/multicard/test_data_parallel.py
           pytest -sv tests/e2e/multicard/test_expert_parallel.py
           pytest -sv tests/e2e/multicard/test_external_launcher.py

tests/e2e/multicard/test_full_graph_mode.py

@@ -29,7 +29,7 @@
 from tests.e2e.model_utils import check_outputs_equal
 
 
-def test_models_distributed_Qwen3_MOE_TP2_WITH_FULLGRAPH():
+def test_models_distributed_Qwen3_MOE_TP2_WITH_FULL_DECODE_ONLY():
     if 'HCCL_OP_EXPANSION_MODE' in os.environ:
         del os.environ['HCCL_OP_EXPANSION_MODE']
     prompts = [
@@ -42,15 +42,64 @@ def test_models_distributed_Qwen3_MOE_TP2_WITH_FULLGRAPH():
                     max_model_len=1024,
                     tensor_parallel_size=2,
                     enforce_eager=False,
-                    compilation_config={"cudagraph_mode":
-                                        "FULL_DECODE_ONLY"}) as runner:
+                    compilation_config={
+                        "cudagraph_mode": "FULL_DECODE_ONLY",
+                        "cudagraph_capture_sizes": [4, 8, 24, 48, 60]
+                    }) as runner:
         vllm_fullgraph_outputs = runner.model.generate(prompts,
                                                        sampling_params)
 
     with VllmRunner(
             model,
             max_model_len=1024,
-            enforce_eager=True,
+            tensor_parallel_size=2,
+            enforce_eager=False,
+    ) as runner:
+        vllm_eager_outputs = runner.model.generate(prompts, sampling_params)
+
+    vllm_fullgraph_outputs_list = []
+    for output in vllm_fullgraph_outputs:
+        vllm_fullgraph_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    vllm_eager_outputs_list = []
+    for output in vllm_eager_outputs:
+        vllm_eager_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    check_outputs_equal(
+        outputs_0_lst=vllm_eager_outputs_list,
+        outputs_1_lst=vllm_fullgraph_outputs_list,
+        name_0="vllm_eager_outputs",
+        name_1="vllm_fullgraph_outputs",
+    )
+
+
+def test_models_distributed_Qwen3_MOE_TP2_WITH_FULL():
+    if 'HCCL_OP_EXPANSION_MODE' in os.environ:
+        del os.environ['HCCL_OP_EXPANSION_MODE']
+    prompts = [
+        "Hello, my name is", "The president of the United States is",
+        "The capital of France is", "The future of AI is"
+    ]
+    model = "Qwen/Qwen3-30B-A3B"
+    sampling_params = SamplingParams(max_tokens=32, temperature=0.0)
+    with VllmRunner(model,
+                    max_model_len=1024,
+                    tensor_parallel_size=2,
+                    enforce_eager=False,
+                    compilation_config={
+                        "cudagraph_mode": "FULL",
+                        "cudagraph_capture_sizes": [4, 8, 24, 48, 60]
+                    }) as runner:
+        vllm_fullgraph_outputs = runner.model.generate(prompts,
+                                                       sampling_params)
+
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            tensor_parallel_size=2,
+            enforce_eager=False,
     ) as runner:
         vllm_eager_outputs = runner.model.generate(prompts, sampling_params)
 

tests/e2e/singlecard/spec_decode_v1/test_v1_mtp_correctness.py

@@ -46,7 +46,7 @@ def mtp_correctness(sampling_config: SamplingParams,
 
     graph_mode_str = "PIECEWISE"
     if graph_mode == CUDAGraphMode.FULL:
-        graph_mode_str = "FULL"
+        graph_mode_str = "FULL_DECODE_ONLY"
 
     with VllmRunner(
             model_name,
@@ -63,7 +63,9 @@ def mtp_correctness(sampling_config: SamplingParams,
             enforce_eager=enforce_eager,
             max_model_len=2000,
             compilation_config=CompilationConfig(
-                cudagraph_mode=graph_mode_str),
+                cudagraph_mode=graph_mode_str,
+                cudagraph_capture_sizes=[12],
+            ),
             additional_config={"ascend_scheduler_config": {
                 "enabled": False
             }}) as spec_llm: