[Compile][Platform] Make PiecewiseBackend pluggable and extendable

[MengqingCao]

Part of #11162

This PR makes PiecewiseBackend pluggable and extensible so that other platforms can easily implement their own segmented graphs.

Other platforms, including oot platforms, that extend to support piecewise static graphs should implement XXXPiecewiseBackend that inherits from AbstractPiecewiseBackend and implement the compilation and graph capture functions in it.

There is already an implementation in vllm-project/vllm-ascend now, and it improves the inference performance significantly

related prs in vllm-project/vllm-ascend:
vllm-project/vllm-ascend#426
vllm-project/vllm-ascend#836

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @MengqingCao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[wangxiyuan]

@youkaichao can you take a look at this? graph mode is a great feature that we hope it can be used widely on different platform

[youkaichao]

LGTM, thanks for the improvement!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @MengqingCao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[MengqingCao]

LGTM, thanks for the improvement!

Thanks! The latest commit resolve the conflict with latest main.

Hope this could be included in v0.9.0 : ) cc @simon-mo @DarkLight1337

[MengqingCao]

Not sure why CI failed.
Test locally with the failed case with this pr, and everything goes well. Could you help to retrigger the CI on entrypoints-test? @youkaichao

(vllm) root@autodl-container-74d51188fa-172f4723:~/admin/vllm# VLLM_USE_MODELSCOPE=True pytest -sv tests/entrypoints/llm/test_generate.py::test_v1_v2_api_consistency_single_prompt_tokens
INFO 05-22 17:04:56 [__init__.py:248] Automatically detected platform cuda.
/root/miniconda/envs/vllm/lib/python3.10/site-packages/pytest_asyncio/plugin.py:208: PytestDeprecationWarning: The configuration option "asyncio_default_fixture_loop_scope" is unset.
The event loop scope for asynchronous fixtures will default to the fixture caching scope. Future versions of pytest-asyncio will default the loop scope for asynchronous fixtures to function scope. Set the default fixture loop scope explicitly in order to avoid unexpected behavior in the future. Valid fixture loop scopes are: "function", "class", "module", "package", "session"

  warnings.warn(PytestDeprecationWarning(_DEFAULT_FIXTURE_LOOP_SCOPE_UNSET))
================================================================== test session starts ==================================================================
platform linux -- Python 3.10.16, pytest-8.3.3, pluggy-1.5.0 -- /root/miniconda/envs/vllm/bin/python
cachedir: .pytest_cache
rootdir: /root/admin/vllm
configfile: pyproject.toml
plugins: anyio-4.6.2.post1, shard-0.1.2, rerunfailures-14.0, forked-1.6.0, asyncio-0.24.0, buildkite-test-collector-0.1.9
asyncio: mode=strict, default_loop_scope=None
collected 4 items                                                                                                                                       
Running 4 items in this shard: tests/entrypoints/llm/test_generate.py::test_v1_v2_api_consistency_single_prompt_tokens[prompt_token_ids0], tests/entrypoints/llm/test_generate.py::test_v1_v2_api_consistency_single_prompt_tokens[prompt_token_ids1], tests/entrypoints/llm/test_generate.py::test_v1_v2_api_consistency_single_prompt_tokens[prompt_token_ids2], tests/entrypoints/llm/test_generate.py::test_v1_v2_api_consistency_single_prompt_tokens[prompt_token_ids3]

tests/entrypoints/llm/test_generate.py::test_v1_v2_api_consistency_single_prompt_tokens[prompt_token_ids0] INFO 05-22 17:05:13 [config.py:787] This model supports multiple tasks: {'classify', 'embed', 'generate', 'reward', 'score'}. Defaulting to 'generate'.
INFO 05-22 17:05:13 [config.py:2112] Chunked prefill is enabled with max_num_batched_tokens=4096.
WARNING 05-22 17:05:13 [cuda.py:87] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be used
INFO 05-22 17:05:15 [core.py:431] Waiting for init message from front-end.
INFO 05-22 17:05:15 [core.py:65] Initializing a V1 LLM engine (v0.9.1.dev44+g23b67b37b.d20250522) with config: model='/root/.cache/modelscope/hub/Qwen/Qwen2___5-0___5B-Instruct', speculative_config=None, tokenizer='/root/.cache/modelscope/hub/Qwen/Qwen2___5-0___5B-Instruct', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config={}, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=32768, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_backend=''), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None), seed=0, served_model_name=/root/.cache/modelscope/hub/Qwen/Qwen2___5-0___5B-Instruct, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=True, use_async_output_proc=False, pooler_config=None, compilation_config={"compile_sizes": [], "inductor_compile_config": {"enable_auto_functionalized_v2": false}, "cudagraph_capture_sizes": [], "max_capture_size": 0}
WARNING 05-22 17:05:16 [utils.py:2664] Methods determine_num_available_blocks,device_config,get_cache_block_size_bytes,initialize_cache not implemented in <vllm.v1.worker.gpu_worker.Worker object at 0x7f5be7af3910>
INFO 05-22 17:05:17 [parallel_state.py:1079] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0, EP rank 0
INFO 05-22 17:05:17 [cuda.py:216] Using Flash Attention backend on V1 engine.
WARNING 05-22 17:05:17 [topk_topp_sampler.py:58] FlashInfer is not available. Falling back to the PyTorch-native implementation of top-p & top-k sampling. For the best performance, please install FlashInfer.
INFO 05-22 17:05:17 [gpu_model_runner.py:1514] Starting to load model /root/.cache/modelscope/hub/Qwen/Qwen2___5-0___5B-Instruct...
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.47it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.46it/s]

INFO 05-22 17:05:17 [default_loader.py:279] Loading weights took 0.48 seconds
INFO 05-22 17:05:18 [gpu_model_runner.py:1532] Model loading took 0.9269 GiB and 0.657718 seconds
INFO 05-22 17:05:19 [kv_cache_utils.py:639] GPU KV cache size: 90,848 tokens
INFO 05-22 17:05:19 [kv_cache_utils.py:642] Maximum concurrency for 32,768 tokens per request: 2.77x
INFO 05-22 17:05:19 [core.py:167] init engine (profile, create kv cache, warmup model) took 0.94 seconds
Adding requests: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 355.51it/s]
Processed prompts: 100%|██████████████████████████████████████████████| 1/1 [00:00<00:00,  3.43it/s, est. speed input: 3.44 toks/s, output: 54.98 toks/s]
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2158.67it/s]
Processed prompts: 100%|██████████████████████████████████████████████| 1/1 [00:00<00:00,  3.77it/s, est. speed input: 3.77 toks/s, output: 60.39 toks/s]
PASSED
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2353.71it/s]
Processed prompts: 100%|██████████████████████████████████████████████| 1/1 [00:00<00:00,  3.69it/s, est. speed input: 7.40 toks/s, output: 59.19 toks/s]
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2327.58it/s]
Processed prompts: 100%|██████████████████████████████████████████████| 1/1 [00:00<00:00,  4.70it/s, est. speed input: 9.41 toks/s, output: 75.23 toks/s]
PASSED
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2398.12it/s]
Processed prompts: 100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  4.58it/s, est. speed input: 13.79 toks/s, output: 73.55 toks/s]
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 1901.32it/s]
Processed prompts: 100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  4.44it/s, est. speed input: 13.35 toks/s, output: 71.17 toks/s]
PASSED
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2351.07it/s]
Processed prompts: 100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  4.58it/s, est. speed input: 18.38 toks/s, output: 73.52 toks/s]
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2321.14it/s]
Processed prompts: 100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  4.14it/s, est. speed input: 16.59 toks/s, output: 66.34 toks/s]
PASSED

================================================================== 4 passed in 21.76s ===================================================================

[DarkLight1337]

These tests are failing on main. Let's wait for them to be fixed first

[MengqingCao]

These tests are failing on main. Let's wait for them to be fixed first

Sure, thanks!

[wangxiyuan]

Hope this feature can be landed in the coming release. @MengqingCao Please rebase once the CI problem is fixed.

[DarkLight1337]

CI has been fixed so please merge from main

[wangxiyuan]

the new commit is just for rebase from main, nothing changed

[MengqingCao]

Thanks @DarkLight1337 for the ci fix and @wangxiyuan for the rebase. Now all the CI is green except for the canceled, I think it's ready for merge