目录
tsingmicro

feat(txda): support cudagraph capture and replay (#649)

PR Category

Vendor

PR Type

New Features

Description

Enable cudagraph (FULL_DECODE_ONLY mode) capture and replay on TXDA (tsingmicro) devices. Previously the tsingmicro triton launcher rejected kernel launches inside the capture scope (txLaunchKernel err=0x47000003) and the device-guard stream diversion corrupted captures, so graph mode was effectively unusable on TXDA.

(none)

Changes

  • vllm_fl/compilation/graph.py: use TXDAGraph on txda; txda-only warmup before capture; apply capture_guard_context() and capture_error_mode="relaxed" at the capture point (non-txda paths byte-identical to upstream)

  • vllm_fl/dispatch/backends/vendor/txda/patch.py (new): stream/current_stream patches, device-guard neutralization, public capture_guard_context()

  • vllm_fl/dispatch/backends/vendor/txda/impl/fused_moe.py, impl/normalization.py (new), register_ops.py, __init__.py

    • vllm_fl/platform.py: enable static graph mode for tsingmicro
    • vllm_fl/worker/worker.py: register graph-capture hook for txda

    Testing

    e2e verified on TXDA (tsingmicro): MiniCPM5-1B, TP=1, --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}', TORCH_COMPILE_DISABLE=1

  • cudagraph capture 1/1 succeeds (Capturing CUDA graphs (decode, FULL): 100%)

  • graph replay output is byte-identical to enforce-eager baseline (same request, temperature=0)

  • runs verified 3/3 in a clean container; non-txda platforms unaffected (nullcontext + empty kwargs path unchanged)

    Checklist

    • I have run the existing tests and they pass
    • I have added tests for my changes (if applicable)
    • I have updated the documentation (if applicable)
1天前181次提交

vllm-plugin-FL

vllm-plugin-FL is a plugin for the vLLM inference/serving framework, built on FlagOS’s unified multi-chip backend — including the unified operator library FlagGems and the unified communication library FlagCX. It extends vLLM’s capabilities and performance across diverse hardware environments. Without changing vLLM’s original interfaces or usage patterns, the same command can run model inference/serving on different chips.

Adaptation Gate

The adaptation gate validates text, image, and mixed text-image requests across the required models in both eager and graph modes. See the adaptation gate cases for the test matrix, scenarios, and usage.

Version Compatibility

vllm-plugin-FL Branch Community vLLM Version
release/0.2 v0.20.2
main v0.24.0

Supported Models and Chips

In theory, vllm-plugin-FL can support all models available in vLLM, as long as no unsupported operators are involved. The tables below summarize the current support status of end-to-end verified models and chips, including both fully supported and in-progress (“Merging”) entries.

Supported Models

Model Status Reference
Qwen3.5-397B-A17B Supported example
Qwen3-Next-80B-A3B Supported example
Qwen3-4B Supported example
MiniCPM-o 4.5 Supported example
GLM-5 Supported example
Qwen3.5-35B-A3B Supported example
BAAI/bge-m3 Supported implementation
MiniMax-M2.7 Supported implementation
MiniMax-M3.1-Flash Supported example / deployment guide
HY4-Preview Merging (NVIDIA) installation

Supported Chips

Chip Vendor Status Reference
NVIDIA Supported -
Ascend Supported -
MetaX Supported -
T-Head Supported -
Iluvatar Supported -
Tsingmicro Supported -
Moore Threads Supported -
Hygon Supported -
Sunrise Supported -
ARM64 CPU Supported -

Quick Start

Setup

  1. Install vLLM

    For NVIDIA GPUs, install vLLM from the official v0.24.0 release (optional if the correct version is already installed):

    pip install vllm==0.24.0

    For non-NVIDIA chips, install vLLM from source with the empty device target:

    git clone -b v0.24.0 https://github.com/vllm-project/vllm.git
    cd vllm
    VLLM_TARGET_DEVICE=empty pip install -v --no-build-isolation --no-deps .
  2. Install vllm-plugin-FL

    2.1 Clone the repository:

    git clone https://github.com/flagos-ai/vllm-plugin-FL

    2.2 Install

    cd vllm-plugin-FL
    pip install --no-build-isolation .
    # or editable install
    pip install --no-build-isolation -e .

    For CUDA-like devices, including CUDA and HIP/ROCm environments that use PyTorch’s CUDA dispatch key, build the plugin native extension by setting VLLM_VENDOR=cuda during installation:

    cd vllm-plugin-FL
    VLLM_VENDOR=cuda pip install --no-build-isolation .
    # or editable install
    VLLM_VENDOR=cuda pip install --no-build-isolation -e .

    This builds and installs vllm_fl._C, which provides native C++ support required by some graph/custom-op paths, especially when vLLM is installed with VLLM_TARGET_DEVICE=empty.

    If VLLM_VENDOR is not set, vllm-plugin-FL is installed as a Python-only plugin and the native extension is skipped.

  3. Install FlagGems

    3.1 Install Build Dependencies

    pip install -U scikit-build-core==0.11 pybind11 ninja cmake

    3.2 Install FlagGems

    git clone -b v5.3.4 https://github.com/flagos-ai/FlagGems
    cd FlagGems
    pip install --no-build-isolation .
    # or editable install
    pip install --no-build-isolation -e .

Runtime compatibility hooks

The plugin installs runtime compatibility hooks through vLLM’s plugin entry points without modifying the installed vLLM package. Model-specific config and model registrations are loaded only for their corresponding architectures.

Operator adapters use the plugin dispatch manager, so backend selection, fallback, per-op policy, operator-list recording, and I/O diagnostics continue to follow the common FlagOS controls.

On Iluvatar with vLLM 0.24, the FlagOS GPU runner binds all attention caches within each decoder layer, including separate main and index caches. Binding preserves layer order, cache order within a layer, tensor identity and shared aliases. This cache initialization fix does not enable graph capture or establish model or operator support; validate those separately in the target environment.

Operator profiling

See the operator profiling tools for a reproducible native-vLLM and plugin comparison workflow.

  1. (Optional) Install FlagCX

    4.1 Clone the repository:

    git clone -b v0.13.0 https://github.com/flagos-ai/FlagCX.git
    cd FlagCX
    git submodule update --init --recursive

    4.2 Build the library with different flags targeting to different platforms:

    make USE_NVIDIA=1

    4.3 Set environment

    export FLAGCX_PATH="$PWD"

    4.4 Installation FlagCX

    cd plugin/torch/
    FLAGCX_ADAPTOR=[xxx] pip install . --no-build-isolation
    # or editable install
    FLAGCX_ADAPTOR=[xxx] pip install -e . --no-build-isolation

    Note: [xxx] should be selected according to the current platform, e.g., nvidia, ascend, etc.

If there are multiple plugins in the current environment, you can specify use vllm-plugin-fl via VLLM_PLUGINS=’fl’.

Additional Steps for Ascend

  1. Install FlagTree

    RES="--index-url=https://resource.flagos.net/repository/flagos-pypi-hosted/simple --trusted-host=https://resource.flagos.net"
    python3 -m pip install flagtree==0.6.1rc1+ascend3.5 $RES

    For other chips, please refer to FlagTree for the corresponding version (e.g., flagtree==0.6.1+iluvatar3.6, flagtree==0.6.1+metax3.6, etc.).

  2. Set required environment variable

    export TRITON_ALL_BLOCKS_PARALLEL=1
  3. Select eager or full-decode graph execution

    For eager execution, add enforce_eager=True to the LLM constructor or pass --enforce-eager on the command line. For FULL_DECODE_ONLY, use:

    export VLLM_PLUGINS=fl
    export USE_FLAGGEMS=0
    
    vllm serve Qwen/Qwen3-0.6B \
        --dtype bfloat16 --tensor-parallel-size 1 \
        --max-model-len 2048 --max-num-batched-tokens 1024 --max-num-seqs 4 \
        --gpu-memory-utilization 0.3 \
        --no-enable-prefix-caching --no-async-scheduling \
        --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,2,4]}'

    This path uses the eager FX backend with NPU graph capture/replay. The standard Python-only FL installation above is sufficient. Use a compatible FlagGems/FlagTree installation when setting USE_FLAGGEMS=1.

HY4 on NVIDIA

HY4 is paired with this branch’s vLLM 0.24.0 environment and the shared FlagGems-vllm integration. Its MLA query concatenation and graph-safe paged MQA logits APIs are now merged into upstream FlagGems-vllm. The shared CI dependency now includes both APIs. Install the same pinned upstream source using the integration instructions above; no separate HY4 upgrade is needed.

HY4 uses the public flaggems_vllm sparse indexer, cache and attention APIs; query quantization uses general FlagGems. No model-specific dependency extra or fork revision is required. The shared fused-operator rollout switch and its default behavior are described in the integration documentation. Runtime selection checks callable APIs, device kernels and the shared dispatch policy before allocating model weights. A selected operator’s execution error propagates without retrying another implementation.

Use hy4_safetensors to load the preview checkpoint with expert slicing. The portable path requires BF16 KV (auto with a BF16 model, or bfloat16); FP8 KV requires the complete native sparse/cache/prefill path. Shared indexers require pipeline parallel size 1. The manual checkpoint smoke configuration requires 16 NVIDIA GPUs; the single-GPU dummy smoke checks model integration. HY4-Preview currently targets NVIDIA H100, vLLM 0.24.0, the MXFP8 preview checkpoint, BF16 activations and BF16 portable KV. Other chips, multimodal serving and portable FP8 KV are not validated. Sparse indexer providers use explicit entry points while preserving the real CUDA platform identity. The portable sparse indexer checks the selected paged MQA logits provider’s graph capability before loading weights. Upstream FlagGems-vllm exposes fp8_fp4_paged_mqa_logits_graph for static-grid CUDA Graph capture, reading lengths and page mappings on the GPU on every replay. The shared dispatch selects that public entry point when available; older operator builds without the capability require --enforce-eager.

With the graph-safe operator build, use FULL_DECODE_ONLY with compilation mode NONE, for example:

--compilation-config '{"mode":"NONE","cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,2,4],"max_cudagraph_capture_size":4}'

The shared CachedOp prewarming from vllm-plugin-FL #552 is used after model loading. Whole-model Dynamo compilation remains limited by the current FlagGems libentry implementation; CUDA Graph mode NONE and eager execution remain available.

Run a Task

Offline Batched Inference

With vLLM and vLLM-fl installed, you can start generating texts for list of input prompts (i.e. offline batch inferencing). See the example script: offline_inference. Or use blow python script directly.

from vllm import LLM, SamplingParams


if __name__ == "__main__":
    prompts = [
        "Hello, my name is",
    ]
    # Create a sampling params object.
    sampling_params = SamplingParams(max_tokens=10, temperature=0.0)
    # Create an LLM.
    llm = LLM(model="Qwen/Qwen3-4B", max_num_batched_tokens=16384, max_num_seqs=2048)
    # Generate texts from the prompts.
    outputs = llm.generate(prompts, sampling_params)
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

Advanced use

For dispatch environment variable usage, see environment variables usage.

Using Cuda Communication library

If you want to use the original Cuda Communication, you can unset the following environment variables.

unset FLAGCX_PATH

Using native CUDA operators

If you want to use the original CUDA operators, you can set the following environment variables.

export USE_FLAGGEMS=0
关于

A vLLM plugin built on the FlagOS unified multi-chip backend.

48.5 MB
邀请码
    Gitlink(确实开源)
  • 加入我们
  • 官网邮箱:gitlink@ccf.org.cn
  • QQ群
  • QQ群
  • 公众号
  • 公众号

版权所有:中国计算机学会技术支持:开源发展技术委员会
京ICP备13000930号-9 京公网安备 11010802047560号