fix activation

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
remove wheel build dependency on extensions
2026-02-19 11:09:03 -05:00 · 2026-02-18 23:57:57 -05:00 · 2026-02-18 22:54:18 -05:00 · 2026-02-18 15:52:50 -08:00 · 2026-02-18 23:10:19 +00:00 · 2026-02-18 15:03:17 -08:00
40 changed files with 1438 additions and 636 deletions
@@ -8,6 +8,21 @@ steps:
  - group: "Build Python wheels"
    key: "build-wheels"
    steps:
+      - label: "Build wheel - aarch64 - CUDA 12.8"
+        depends_on: ~
+        id: build-wheel-arm64-cuda-12-8
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        commands:
+          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
+          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+        env:
+          DOCKER_BUILDKIT: "1"
+
      - label: "Build wheel - aarch64 - CUDA 12.9"
        depends_on: ~
        id: build-wheel-arm64-cuda-12-9
@@ -51,6 +66,19 @@ steps:
        env:
          DOCKER_BUILDKIT: "1"

+      - label: "Build wheel - x86_64 - CUDA 12.8"
+        depends_on: ~
+        id: build-wheel-x86-cuda-12-8
+        agents:
+          queue: cpu_queue_postmerge
+        commands:
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
+        env:
+          DOCKER_BUILDKIT: "1"
+
      - label: "Build wheel - x86_64 - CUDA 12.9"
        depends_on: ~
        id: build-wheel-x86-cuda-12-9
@@ -19,7 +19,7 @@ else()
  FetchContent_Declare(
        flashmla
        GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
+        GIT_TAG f46984b4caf19409bf4ab94a540341918dd6a1a3
        GIT_PROGRESS TRUE
        CONFIGURE_COMMAND ""
        BUILD_COMMAND ""
@@ -13,8 +13,13 @@ struct alignas(32) u32x8_t {
  uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
 };

+// 256-bit vector loads require both sm_100+ and CUDA 12.9+
+#define VLLM_SUPPORTS_256BIT_VECTORS                  \
+  (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
+   __CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 9)
+
 __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+#if VLLM_SUPPORTS_256BIT_VECTORS
  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
               : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
                 "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
@@ -35,7 +40,7 @@ __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
 }

 __device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+#if VLLM_SUPPORTS_256BIT_VECTORS
  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
               :
               : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
@@ -418,10 +418,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    fi && \
    python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38

-# Copy extension wheels from extensions-build stage for later use
-COPY --from=extensions-build /tmp/deepgemm/dist /tmp/deepgemm/dist
-COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist
-
 # Check the size of the wheel if RUN_WHEEL_CHECK is true
 COPY .buildkite/check-wheel-size.py check-wheel-size.py
 # sync the default value with .buildkite/check-wheel-size.py
@@ -660,9 +656,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 . /etc/environment && \
 uv pip list

-# Install deepgemm wheel that has been built in the `build` stage
+# Install deepgemm wheel that has been built in the `extensions-build` stage
 RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
+    --mount=type=bind,from=extensions-build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
    sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \
              uv pip install --system /tmp/deepgemm/dist/*.whl; \
           else \
@@ -672,8 +668,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
 ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

-# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage
-RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
+# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `extensions-build` stage
+RUN --mount=type=bind,from=extensions-build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
    --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system ep_kernels/dist/*.whl --verbose \
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -36,12 +36,12 @@ th:not(:first-child) {
 }
 </style>

-| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
+| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](speculative_decoding/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
 |---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
 | [CP](../configuration/optimization.md#chunked-prefill) | ✅ | | | | | | | | | | | | | | |
 | [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
 | [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
-| [SD](spec_decode/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
+| [SD](speculative_decoding/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
 | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
 | [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ✅ | ✅ | ✅ | | | | | | | | |
@@ -64,7 +64,7 @@ th:not(:first-child) {
 | [CP](../configuration/optimization.md#chunked-prefill)                                     | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
 | [APC](automatic_prefix_caching.md)                        | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
 | [LoRA](lora.md)                                           | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
-| [SD](spec_decode/README.md)                        | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | ✅        |
+| [SD](speculative_decoding/README.md)                        | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | ✅        |
 | CUDA graph                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | [❌](https://github.com/vllm-project/vllm/issues/26970)        |
 | [pooling](../models/pooling_models.md)                    | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr>       | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❌     | ✅        |
@@ -1,330 +0,0 @@
-# Speculative Decoding
-
-!!! warning
-    Currently, speculative decoding in vLLM is not compatible with pipeline parallelism.
-
-This document shows how to use [Speculative Decoding](https://x.com/karpathy/status/1697318534555336961) with vLLM.
-Speculative decoding is a technique which improves inter-token latency in memory-bound LLM inference.
-
-!!! tip
-    To train your own draft models for speculative decoding, see [Speculators](speculators.md), a library for training draft models that integrates seamlessly with vLLM.
-
-## Speculating with a draft model
-
-The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
-
-!!! warning
-    In vllm v0.10.0, speculative decoding with a draft model is not supported.
-    If you use the following code, you will get a `NotImplementedError`.
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="facebook/opt-6.7b",
-        tensor_parallel_size=1,
-        speculative_config={
-            "model": "facebook/opt-125m",
-            "num_speculative_tokens": 5,
-        },
-    )
-    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}")
-    ```
-
-To perform the same with an online mode launch the server:
-
-```bash
-vllm serve facebook/opt-6.7b \
-    --host 0.0.0.0 \
-    --port 8000 \
-    --seed 42 \
-    -tp 1 \
-    --gpu_memory_utilization 0.8 \
-    --speculative_config '{"model": "facebook/opt-125m", "num_speculative_tokens": 5}'
-```
-
-!!! warning
-    Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated now.
-
-Then use a client:
-
-??? code
-
-    ```python
-    from openai import OpenAI
-
-    # Modify OpenAI's API key and API base to use vLLM's API server.
-    openai_api_key = "EMPTY"
-    openai_api_base = "http://localhost:8000/v1"
-
-    client = OpenAI(
-        # defaults to os.environ.get("OPENAI_API_KEY")
-        api_key=openai_api_key,
-        base_url=openai_api_base,
-    )
-
-    models = client.models.list()
-    model = models.data[0].id
-
-    # Completion API
-    stream = False
-    completion = client.completions.create(
-        model=model,
-        prompt="The future of AI is",
-        echo=False,
-        n=1,
-        stream=stream,
-    )
-
-    print("Completion results:")
-    if stream:
-        for c in completion:
-            print(c)
-    else:
-        print(completion)
-    ```
-
-## Speculating by matching n-grams in the prompt
-
-The following code configures vLLM to use speculative decoding where proposals are generated by
-matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="facebook/opt-6.7b",
-        tensor_parallel_size=1,
-        speculative_config={
-            "method": "ngram",
-            "num_speculative_tokens": 5,
-            "prompt_lookup_max": 4,
-        },
-    )
-    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}")
-    ```
-
-## Speculating using Suffix Decoding
-
-The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
-
-Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
-
-Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
-
-!!! tip "Install Arctic Inference"
-    Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
-
-!!! tip "Suffix Decoding Speculative Tokens"
-    Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="facebook/opt-6.7b",
-        tensor_parallel_size=1,
-        speculative_config={
-            "method": "suffix",
-            "num_speculative_tokens": 32,
-        },
-    )
-    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}")
-    ```
-
-## Speculating using MLP speculators
-
-The following code configures vLLM to use speculative decoding where proposals are generated by
-draft models that condition draft predictions on both context vectors and sampled tokens.
-For more information see [this blog](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) or
-[this technical report](https://arxiv.org/abs/2404.19124).
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="meta-llama/Meta-Llama-3.1-70B-Instruct",
-        tensor_parallel_size=4,
-        speculative_config={
-            "model": "ibm-ai-platform/llama3-70b-accelerator",
-            "draft_tensor_parallel_size": 1,
-        },
-    )
-    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}")
-    ```
-
-Note that these speculative models currently need to be run without tensor parallelism, although
-it is possible to run the main model using tensor parallelism (see example above). Since the
-speculative models are relatively small, we still see significant speedups. However, this
-limitation will be fixed in a future release.
-
-A variety of speculative models of this type are available on HF hub:
-
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
-
-## Speculating using EAGLE based draft models
-
-The following code configures vLLM to use speculative decoding where proposals are generated by
-an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="meta-llama/Meta-Llama-3-8B-Instruct",
-        tensor_parallel_size=4,
-        speculative_config={
-            "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
-            "draft_tensor_parallel_size": 1,
-            "num_speculative_tokens": 2,
-            "method": "eagle",
-        },
-    )
-
-    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}")
-
-    ```
-
-A few important things to consider when using the EAGLE based draft models:
-
-1. The EAGLE draft models available in the [HF repository for EAGLE models](https://huggingface.co/yuhuili) should
-   be able to be loaded and used directly by vLLM after <https://github.com/vllm-project/vllm/pull/12304>.
-   If you are using vllm version before <https://github.com/vllm-project/vllm/pull/12304>, please use the
-   [script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model,
-   and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`. If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.
-
-2. The EAGLE based draft models need to be run without tensor parallelism
-   (i.e. draft_tensor_parallel_size is set to 1 in `speculative_config`), although
-   it is possible to run the main model using tensor parallelism (see example above).
-
-3. When using EAGLE-based speculators with vLLM, the observed speedup is lower than what is
-   reported in the reference implementation [here](https://github.com/SafeAILab/EAGLE). This issue is under
-   investigation and tracked here: <https://github.com/vllm-project/vllm/issues/9565>.
-
-4. When using EAGLE-3 based draft model, option "method" must be set to "eagle3".
-   That is, to specify `"method": "eagle3"` in `speculative_config`.
-
-A variety of EAGLE draft models are available on the Hugging Face hub:
-
-| Base Model                                                           | EAGLE on Hugging Face                     | # EAGLE Parameters |
-|---------------------------------------------------------------------|-------------------------------------------|--------------------|
-| Vicuna-7B-v1.3                                                       | yuhuili/EAGLE-Vicuna-7B-v1.3             | 0.24B              |
-| Vicuna-13B-v1.3                                                      | yuhuili/EAGLE-Vicuna-13B-v1.3            | 0.37B              |
-| Vicuna-33B-v1.3                                                      | yuhuili/EAGLE-Vicuna-33B-v1.3            | 0.56B              |
-| LLaMA2-Chat 7B                                                       | yuhuili/EAGLE-llama2-chat-7B             | 0.24B              |
-| LLaMA2-Chat 13B                                                      | yuhuili/EAGLE-llama2-chat-13B            | 0.37B              |
-| LLaMA2-Chat 70B                                                      | yuhuili/EAGLE-llama2-chat-70B            | 0.99B              |
-| Mixtral-8x7B-Instruct-v0.1                                           | yuhuili/EAGLE-mixtral-instruct-8x7B      | 0.28B              |
-| LLaMA3-Instruct 8B                                                   | yuhuili/EAGLE-LLaMA3-Instruct-8B         | 0.25B              |
-| LLaMA3-Instruct 70B                                                  | yuhuili/EAGLE-LLaMA3-Instruct-70B        | 0.99B              |
-| Qwen2-7B-Instruct                                                    | yuhuili/EAGLE-Qwen2-7B-Instruct          | 0.26B              |
-| Qwen2-72B-Instruct                                                   | yuhuili/EAGLE-Qwen2-72B-Instruct         | 1.05B              |
-
-## Lossless guarantees of Speculative Decoding
-
-In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
-speculative decoding, breaking down the guarantees into three key areas:
-
-1. **Theoretical Losslessness**
-   \- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
-   cause slight variations in output distributions, as discussed
-   in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
-
-2. **Algorithmic Losslessness**
-   \- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
-
-    > - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
-    >   distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
-    > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
-    >   without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
-    >   provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../tests/spec_decode/e2e).
-    >   verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
-
-3. **vLLM Logprob Stability**
-   \- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
-   same request across runs. For more details, see the FAQ section
-   titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
-
-While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
-can occur due to following factors:
-
- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
-  due to non-deterministic behavior in batched operations or numerical instability.
-
-For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
-
-## Resources for vLLM contributors
-
- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
@@ -0,0 +1,62 @@
+# Speculative Decoding
+
+This document shows how to use [Speculative Decoding](https://arxiv.org/pdf/2302.01318) with vLLM to reduce inter-token latency under medium-to-low QPS (query per second), memory-bound workloads.
+
+To train your own draft models for optimized speculative decoding, see [vllm-project/speculators](speculators.md) for seamless training and integration with vLLM.
+
+## vLLM Speculation Methods
+
+vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, draft models, and mlp provide the best latency reduction, while simpler methods such as n-gram and and suffix decoding provide modest speedups without increasing workload during peak traffic.
+
+- [EAGLE](eagle.md)
+- [Draft Model](draft_model.md)
+- [Multi-Layer Perceptron](mlp.md)
+- [N-Gram](n_gram.md)
+- [Suffix Decoding](suffix.md)
+
+## Lossless guarantees of Speculative Decoding
+
+In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
+speculative decoding, breaking down the guarantees into three key areas:
+
+1. **Theoretical Losslessness**
+   \- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
+   cause slight variations in output distributions, as discussed
+   in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
+
+2. **Algorithmic Losslessness**
+   \- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
+
+    > - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
+    >   distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
+    > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
+    >   without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
+    >   provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
+    >   verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
+
+3. **vLLM Logprob Stability**
+   \- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
+   same request across runs. For more details, see the FAQ section
+   titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
+
+While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
+can occur due to following factors:
+
+- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
+- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
+  due to non-deterministic behavior in batched operations or numerical instability.
+
+For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
+
+## Known Feature Incompatibility
+
+1. Pipeline parallelism is not composible with speculative decoding as of `vllm<=0.15.0`
+2. Speculative decoding with a draft models is not supported in `vllm<=0.10.0`
+
+## Resources for vLLM contributors
+
+- [[vLLM Office Hours #40] Intro to Speculators](https://www.youtube.com/watch?v=2ISAr_JVGLs)
+- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
+- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
+- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
+- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
@@ -0,0 +1,80 @@
+# Draft Models
+
+The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    tensor_parallel_size=1,
+    speculative_config={
+        "model": "Qwen/Qwen3-0.6B",
+        "num_speculative_tokens": 5,
+        "method": "draft_model",
+    },
+)
+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}")
+```
+
+To perform the equivalent launch in online mode, use the following server-side code:
+
+```bash
+vllm serve Qwen/Qwen3-4B-Thinking-2507 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --seed 42 \
+    -tp 1 \
+    --max_model_len 2048 \
+    --gpu_memory_utilization 0.8 \
+    --speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
+```
+
+The code used to request as completions as a client remains unchanged:
+
+??? code
+
+    ```python
+    from openai import OpenAI
+
+    # Modify OpenAI's API key and API base to use vLLM's API server.
+    openai_api_key = "EMPTY"
+    openai_api_base = "http://localhost:8000/v1"
+
+    client = OpenAI(
+        # defaults to os.environ.get("OPENAI_API_KEY")
+        api_key=openai_api_key,
+        base_url=openai_api_base,
+    )
+
+    models = client.models.list()
+    model = models.data[0].id
+
+    # Completion API
+    stream = False
+    completion = client.completions.create(
+        model=model,
+        prompt="The future of AI is",
+        echo=False,
+        n=1,
+        stream=stream,
+    )
+
+    print("Completion results:")
+    if stream:
+        for c in completion:
+            print(c)
+    else:
+        print(completion)
+    ```
+
+!!! warning
+    Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
@@ -0,0 +1,67 @@
+# EAGLE Draft Models
+
+The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
+
+## Eagle Drafter Example
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="meta-llama/Meta-Llama-3-8B-Instruct",
+    tensor_parallel_size=4,
+    speculative_config={
+        "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
+        "draft_tensor_parallel_size": 1,
+        "num_speculative_tokens": 2,
+        "method": "eagle",
+    },
+)
+
+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}")
+```
+
+## Eagle3 Drafter Example
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="meta-llama/Meta-Llama-3-8B-Instruct",
+    tensor_parallel_size=2,
+    speculative_config={
+        "model": "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
+        "draft_tensor_parallel_size": 2,
+        "num_speculative_tokens": 2,
+        "method": "eagle",
+    },
+)
+
+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}")
+```
+
+## Pre-Trained Eagle Draft Models
+
+A variety of EAGLE draft models are available on the Hugging Face hub:
+
+* [RedHatAI/speculator-models](https://huggingface.co/collections/RedHatAI/speculator-models)
+* [yuhuili/models](https://huggingface.co/yuhuili/models?search=eagle)
+
+!!! warning
+    If you are using `vllm<0.7.0`, please use [this script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`.
@@ -0,0 +1,42 @@
+# MLP Draft Models
+
+The following code configures vLLM to use speculative decoding where proposals are generated by draft models that condition draft predictions on both context vectors and sampled tokens. For more information see [The Hitchhiker's Guide to Speculative Decoding](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) and [IBM Research's Technical Report](https://arxiv.org/abs/2404.19124).
+
+## MLP Drafter Example
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
+    tensor_parallel_size=4,
+    speculative_config={
+        "model": "ibm-ai-platform/llama3-70b-accelerator",
+        "draft_tensor_parallel_size": 1,
+        "method": "mlp_speculator",
+    },
+)
+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}")
+```
+
+## Pre-Trained MLP Drafter Models
+
+A variety of speculative models of this type are available on HF hub:
+
+- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
+- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
+- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
+- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
+- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
+- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
+- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
+- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
+- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
@@ -0,0 +1,27 @@
+# N-Gram Speculation
+
+The following code configures vLLM to use speculative decoding where proposals are generated by
+matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    tensor_parallel_size=1,
+    speculative_config={
+        "method": "ngram",
+        "num_speculative_tokens": 5,
+        "prompt_lookup_max": 4,
+    },
+)
+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}")
+```
@@ -1,4 +1,7 @@
-# Speculators
+# vLLM-Project/Speculators
+
+![User Flow Light](../../assets/features/speculative_decoding/speculators-user-flow-light.svg#only-light)
+![User Flow Dark](../../assets/features/speculative_decoding/speculators-user-flow-dark.svg#only-dark)

 [Speculators](https://docs.vllm.ai/projects/speculators/en/latest/) is a library for accelerating LLM inference through speculative decoding, providing efficient draft model training that integrates seamlessly with vLLM to reduce latency and improve throughput.

@@ -0,0 +1,35 @@
+# Suffix Decoding
+
+The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
+
+Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
+
+Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
+
+!!! tip "Install Arctic Inference"
+    Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
+
+!!! tip "Suffix Decoding Speculative Tokens"
+    Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    tensor_parallel_size=1,
+    speculative_config={
+        "method": "suffix",
+        "num_speculative_tokens": 32,
+    },
+)
+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}")
+```
@@ -105,6 +105,10 @@ plugins:
          - https://numpy.org/doc/stable/objects.inv
          - https://pytorch.org/docs/stable/objects.inv
          - https://psutil.readthedocs.io/en/stable/objects.inv
+  - redirects:
+      redirect_maps:
+        features/spec_decode/README.md: features/speculative_decoding/README.md
+        features/spec_decode/speculators.md: features/speculative_decoding/speculators.md

 markdown_extensions:
  - attr_list
@@ -7,6 +7,7 @@ mkdocs-awesome-nav
 mkdocs-glightbox
 mkdocs-git-revision-date-localized-plugin
 mkdocs-minify-plugin
+mkdocs-redirects
 regex
 ruff
 pydantic
@@ -977,9 +977,9 @@ if _is_cuda():
        # FA3 requires CUDA 12.3 or later
        ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
    if envs.VLLM_USE_PRECOMPILED or (
-        CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
+        CUDA_HOME and get_nvcc_cuda_version() >= Version("12.8")
    ):
-        # FlashMLA requires CUDA 12.9 or later
+        # FlashMLA requires CUDA 12.8 or later
        # Optional since this doesn't get built (produce an .so file) when
        # not targeting a hopper system
        ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
@@ -124,6 +124,8 @@ def test_init_weight_transfer_engine_calls_engine():
    if torch.cuda.device_count() < 1:
        pytest.skip("Need at least 1 GPU for this test")

+    # Run in-process so mock.patch works (spawn won't inherit the mock)
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    # Enable insecure serialization to allow pickling functions for collective_rpc
    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

@@ -171,6 +173,8 @@ def test_update_weights_calls_engine():
    if torch.cuda.device_count() < 1:
        pytest.skip("Need at least 1 GPU for this test")

+    # Run in-process so mock.patch works (spawn won't inherit the mock)
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    # Enable insecure serialization to allow pickling functions for collective_rpc
    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

@@ -228,6 +232,8 @@ def test_full_weight_transfer_flow():
    if torch.cuda.device_count() < 1:
        pytest.skip("Need at least 1 GPU for this test")

+    # Run in-process so mock.patch works (spawn won't inherit the mock)
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    # Enable insecure serialization to allow pickling functions for collective_rpc
    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

@@ -398,80 +398,3 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(

    assert w13_converted.shape[0] == num_experts
    assert w2_converted.shape[0] == num_experts
-
-
-def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
-    """Test that flashinfer_fused_moe_blockscale_fp8 handles num_expert_group=None.
-
-    Regression test for https://github.com/vllm-project/vllm/issues/34477
-    MiniMax-M2.1 uses sigmoid scoring with e_score_correction_bias but no
-    grouped top-k, resulting in num_expert_group=None. This triggered a crash
-    in the flashinfer kernel when DeepSeekV3 routing was selected.
-    """
-    if not current_platform.has_device_capability(100):
-        pytest.skip("Test requires SM >= 100 (Blackwell)")
-
-    import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe  # noqa: E501, F401
-    from tests.kernels.quant_utils import native_per_token_group_quant_fp8
-
-    set_random_seed(7)
-    monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
-
-    e = 16  # num_experts (must be divisible by 4)
-    topk = 6  # top_k > 1 triggers DeepSeekV3 routing with sigmoid
-    m, n, k = 10, 4096, 5120
-    block_shape = [128, 128]
-    block_k = block_shape[1]
-
-    with set_current_vllm_config(vllm_config):
-        # Create BF16 hidden states
-        x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
-
-        # Create FP8 block-scale quantized weights
-        w13_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=torch.bfloat16) / 10
-        w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10
-
-        # Quantize weights per-block to FP8
-        w13_fp8_list, w13_scale_list = [], []
-        w2_fp8_list, w2_scale_list = [], []
-        for i in range(e):
-            wq, ws = native_per_token_group_quant_fp8(w13_bf16[i], block_k)
-            w13_fp8_list.append(wq)
-            w13_scale_list.append(ws)
-
-            wq, ws = native_per_token_group_quant_fp8(w2_bf16[i], block_k)
-            w2_fp8_list.append(wq)
-            w2_scale_list.append(ws)
-
-        w13_fp8 = torch.stack(w13_fp8_list)
-        w13_scale = torch.stack(w13_scale_list)
-        w2_fp8 = torch.stack(w2_fp8_list)
-        w2_scale = torch.stack(w2_scale_list)
-
-        # DeepSeekV3 routing uses float32 logits + optional bias
-        routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
-        routing_bias = torch.randn(e, device="cuda", dtype=torch.float32)
-
-        # This should NOT crash with num_expert_group=None
-        output = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
-            routing_logits=routing_logits,
-            routing_bias=routing_bias,
-            x=x,
-            w13_weight=w13_fp8,
-            w13_weight_scale_inv=w13_scale,
-            w2_weight=w2_fp8,
-            w2_weight_scale_inv=w2_scale,
-            global_num_experts=e,
-            top_k=topk,
-            num_expert_group=None,
-            topk_group=None,
-            intermediate_size=n,
-            expert_offset=0,
-            local_num_experts=e,
-            block_shape=block_shape,
-            routing_method_type=RoutingMethodType.DeepSeekV3,
-            routed_scaling=1.0,
-        )
-
-        assert output is not None
-        assert output.shape == (m, k)
@@ -155,9 +155,9 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
    out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)

    if xnorm:
-        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
+        torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
    else:
-        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
+        torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
@@ -177,7 +177,7 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
    ref_out = torch.matmul(A, B.t())
    out = ops.LLMM1(B, A, rows_per_block)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -194,7 +194,7 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
    ref_out = torch.nn.functional.linear(A, B)
    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -213,7 +213,7 @@ def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
    ref_out = torch.nn.functional.linear(A, B, BIAS)
    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -232,7 +232,7 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
    ref_out = torch.nn.functional.linear(A, B, BIAS)
    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("xnorm", [False, True])
@@ -275,4 +275,4 @@ def test_rocm_wvsplitk_fp8_kernel(
        # wider pytrch thresh for large-K & no xnorm
        torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2)
    else:
-        torch.testing.assert_close(out, ref_out, atol=0.01, rtol=0.01)
+        torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
@@ -153,5 +153,5 @@ def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
        # Then check to make sure the submitted lora request
        # and text prompt were zipped together correctly
        engine_args, engine_kwargs = mock_add_request.call_args
+        assert engine_args[1]["prompt"] == AUDIO_PROMPT
        assert engine_kwargs["lora_request"] is None
-        assert engine_kwargs["prompt_text"] == AUDIO_PROMPT
@@ -88,9 +88,8 @@ class Qwen2VLTester:
        # Validate outputs
        for generated, expected in zip(generated_texts, expected_outputs):
            assert expected.startswith(generated), (
-                f"Generated text {generated} doesn't "
+                f"Generated text {generated} doesn't match expected pattern {expected}"
            )
-            f"match expected pattern {expected}"

    def run_beam_search_test(
        self,
@@ -118,11 +117,14 @@ class Qwen2VLTester:
            inputs, beam_search_params, lora_request=lora_request
        )

-        for output_obj, expected_outs in zip(outputs, expected_outputs):
+        for output_obj, expected_texts in zip(outputs, expected_outputs):
            output_texts = [seq.text for seq in output_obj.sequences]
-            assert output_texts == expected_outs, (
-                f"Generated texts {output_texts} do not match expected {expected_outs}"
-            )  # noqa: E501
+
+            for output_text, expected_text in zip(output_texts, expected_texts):
+                # NOTE beam search .text contains the whole text including inputs
+                assert output_text.endswith(expected_text), (
+                    f"Generated {output_text} does not match expected {expected_text}"
+                )


 TEST_IMAGES = [
@@ -151,11 +153,10 @@ EXPECTED_OUTPUTS_VISION_NO_CONNECTOR = [
    "A closeup shot of the Tokyo Skytree with pink flowers in the foreground.",
 ]

-# NOTE - beam search .text contains the whole text
 EXPECTED_BEAM_SEARCH_OUTPUTS = [
    [
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>\n<|im_start|>assistant\nA majestic skyscraper stands",  # noqa: E501
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>\n<|im_start|>assistant\nA majestic tower stands tall",  # noqa: E501
+        "A majestic skyscraper stands",
+        "A majestic tower stands tall",
    ],
 ]

@@ -8,7 +8,7 @@ MODEL = "meta-llama/Llama-3.2-1B-Instruct"


 def test_unsupported_configs():
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
        AsyncEngineArgs(
            model=MODEL,
            speculative_config={
@@ -305,6 +305,13 @@ class SpeculativeConfig:
        # can not be detected, it will be considered as the "draft_model" by
        # default.

+        # infer method from user args
+        if self.method is None:
+            if self.model in ("ngram", "[ngram]"):
+                self.method = "ngram"
+            else:
+                self.method = "draft_model"
+
        if self.method in get_args(MTPModelTypes) and self.method != "mtp":
            logger.warning(
                "method `%s` is deprecated and replaced with mtp.", self.method
@@ -334,13 +341,6 @@ class SpeculativeConfig:
                    "num_speculative_tokens was provided but without speculative model."
                )

-        # Automatically configure the method for ngram when "model" is used
-        # instead of "method"
-        if self.method is None and (
-            self.model is not None and self.model in ("ngram", "[ngram]")
-        ):
-            self.method = "ngram"
-
        if self.method in ("ngram", "[ngram]"):
            # Unified to "ngram" internally
            self.method = "ngram"
@@ -505,6 +505,13 @@ class SpeculativeConfig:
                        )

                if self.speculative_token_tree is None:
+                    if self.num_speculative_tokens is None:
+                        raise ValueError(
+                            "A speculative model was provided, but neither "
+                            "`speculative_token_tree` nor `num_speculative_tokens` "
+                            "was provided"
+                        )
+
                    # Generate chain of tokens.
                    self.speculative_token_tree = str(
                        [(i + 1) * (0,) for i in range(self.num_speculative_tokens)]
@@ -40,6 +40,7 @@ class AnthropicContentBlock(BaseModel):
    source: dict[str, Any] | None = None
    # For tool use/result
    id: str | None = None
+    tool_use_id: str | None = None
    name: str | None = None
    input: dict[str, Any] | None = None
    content: str | list[dict[str, Any]] | None = None
@@ -139,7 +139,7 @@ class AnthropicServingMessages(OpenAIServingChat):
                            openai_messages.append(
                                {
                                    "role": "tool",
-                                    "tool_call_id": block.id or "",
+                                    "tool_call_id": block.tool_use_id or "",
                                    "content": str(block.content)
                                    if block.content
                                    else "",
@@ -542,51 +542,31 @@ class LLM:
        outputs = self._run_engine(use_tqdm=use_tqdm)
        return self.engine_class.validate_outputs(outputs, RequestOutput)

-    def _resolve_lora_reqs(
-        self,
-        prompts: Sequence[ProcessorInputs],
-        lora_request: Sequence[LoRARequest | None] | LoRARequest | None,
-    ):
-        lora_config = self.llm_engine.vllm_config.lora_config
-        seq_lora_requests = self._lora_request_to_seq(lora_request, len(prompts))
-
-        if (
-            lora_config is None
-            or not self.model_config.is_multimodal_model
-            or (lora_config and lora_config.default_mm_loras is None)
-        ):
-            return seq_lora_requests
-
-        return [
-            self._resolve_single_prompt_mm_lora(
-                prompt,
-                lora_req,
-                lora_config.default_mm_loras,
-            )
-            for prompt, lora_req in zip(prompts, seq_lora_requests)
-        ]
-
-    def _resolve_single_prompt_mm_lora(
+    def _resolve_mm_lora(
        self,
        prompt: ProcessorInputs,
        lora_request: LoRARequest | None,
-        default_mm_loras: dict[str, str] | None,
-    ):
-        if not default_mm_loras or prompt["type"] != "multimodal":
+    ) -> LoRARequest | None:
+        if prompt["type"] != "multimodal":
+            return lora_request
+
+        lora_config = self.llm_engine.vllm_config.lora_config
+        default_mm_loras = None if lora_config is None else lora_config.default_mm_loras
+        if not default_mm_loras:
            return lora_request

        prompt_modalities = prompt["mm_placeholders"].keys()
        intersection = set(prompt_modalities).intersection(default_mm_loras.keys())
        if not intersection:
            return lora_request
+
        if len(intersection) > 1:
            # TODO: Would be nice to be able to have multiple loras per prompt
            logger.warning(
-                "Multiple modality specific loras were registered and would be"
-                " used by a single prompt consuming several modalities; "
-                " currently we only support one lora per request; as such,"
-                " lora(s) registered with modalities: %s"
-                " will be skipped",
+                "Multiple modality specific loras were registered and would be "
+                "used by a single prompt consuming several modalities; "
+                "currently we only support one lora per request; as such, "
+                "lora(s) registered with modalities: %s will be skipped",
                intersection,
            )
            return lora_request
@@ -1915,7 +1895,10 @@ class LLM:
                request_id = self._add_request(
                    prompt,
                    params[i],
-                    lora_request=None if lora_requests is None else lora_requests[i],
+                    lora_request=self._resolve_mm_lora(
+                        prompt,
+                        None if lora_requests is None else lora_requests[i],
+                    ),
                    priority=0 if priorities is None else priorities[i],
                )
                added_request_ids.append(request_id)
@@ -8,11 +8,7 @@ import torch

 import vllm.envs as envs
 from vllm.config import ParallelConfig
-from vllm.distributed import (
-    get_dp_group,
-    get_pcp_group,
-    get_tensor_model_parallel_rank,
-)
+from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -126,20 +122,31 @@ class RoutingMethodType(IntEnum):


 def get_routing_method_type(
-    scoring_func: str, top_k: int, renormalize: bool
+    scoring_func: str,
+    top_k: int,
+    renormalize: bool,
+    num_expert_group: int | None,
+    has_e_score_bias: bool,
 ) -> RoutingMethodType:
+    if has_e_score_bias:
+        if (num_expert_group or 0) > 0 and scoring_func == "sigmoid":
+            return RoutingMethodType.DeepSeekV3
+        else:
+            return RoutingMethodType.Unspecified
+
    if scoring_func == "sigmoid":
        if top_k == 1:
            return RoutingMethodType.Llama4
        else:
-            return RoutingMethodType.DeepSeekV3
-    elif scoring_func == "softmax":
+            return RoutingMethodType.Unspecified
+
+    if scoring_func == "softmax":
        if renormalize:
            return RoutingMethodType.Renormalize
        else:
            return RoutingMethodType.Default
-    else:
-        return RoutingMethodType.Unspecified
+
+    return RoutingMethodType.Unspecified


@dataclass
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
    disable_inplace,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    get_marlin_input_dtype,
    marlin_make_workspace_new,
    marlin_moe_intermediate_size,
    marlin_quant_input,
@@ -550,6 +551,8 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):
        self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
        self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
        self.is_k_full = is_k_full
+        self.input_dtype = get_marlin_input_dtype()
+
        super().__init__(
            moe_config=moe_config,
            quant_config=quant_config,
@@ -736,6 +739,7 @@ class MarlinExperts(MarlinExpertsBase):
            sort_indices1=self.w13_g_idx_sort_indices,
            sort_indices2=self.w2_g_idx_sort_indices,
            is_k_full=self.is_k_full,
+            input_dtype=self.input_dtype,
        )

    def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
@@ -165,6 +165,8 @@ class FusedTopKBiasRouter(BaseRouter):
            scoring_func=self.scoring_func,
            top_k=self.top_k,
            renormalize=self.renormalize,
+            num_expert_group=None,
+            has_e_score_bias=True,
        )

    def _compute_routing(
@@ -142,6 +142,8 @@ class FusedTopKRouter(BaseRouter):
            scoring_func=self.scoring_func,
            top_k=self.top_k,
            renormalize=self.renormalize,
+            num_expert_group=None,
+            has_e_score_bias=False,
        )

    def _compute_routing(
@@ -13,7 +13,10 @@ from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.batch_invariant import (
    vllm_is_batch_invariant,
 )
-from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+from vllm.model_executor.layers.fused_moe.config import (
+    RoutingMethodType,
+    get_routing_method_type,
+)
 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
    rocm_aiter_grouped_topk,
 )
@@ -277,16 +280,15 @@ class GroupedTopKRouter(BaseRouter):
        self.e_score_correction_bias = e_score_correction_bias
        self.num_fused_shared_experts = num_fused_shared_experts

-        if scoring_func == "sigmoid":
-            self._routing_method_type = RoutingMethodType.DeepSeekV3
-        else:
-            # NOTE: this prohibits the FLASHINFER_TRTLLM kernels from
-            # being selected, since they only support DeepSeek-style.
-            self._routing_method_type = RoutingMethodType.Unspecified
-
    @property
    def routing_method_type(self) -> RoutingMethodType:
-        return self._routing_method_type
+        return get_routing_method_type(
+            scoring_func=self.scoring_func,
+            top_k=self.top_k,
+            renormalize=self.renormalize,
+            num_expert_group=self.num_expert_group,
+            has_e_score_bias=self.e_score_correction_bias is not None,
+        )

    def _compute_routing(
        self,
@@ -16,6 +16,9 @@ from vllm.model_executor.layers.fused_moe import (
    MoEActivation,
 )
 from vllm.model_executor.layers.fused_moe import modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    maybe_make_prepare_finalize,
+)
 from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
    mxfp4_mxfp8_moe_quant_config,
@@ -25,7 +28,6 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
    BatchedMarlinExperts,
    MarlinExperts,
-    fused_marlin_moe,
 )
 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
    OAITritonExperts,
@@ -52,7 +54,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
-from vllm.scalar_type import scalar_types
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.utils.import_utils import has_triton_kernels
 from vllm.utils.math_utils import round_up
@@ -219,7 +220,6 @@ class Mxfp4Config(QuantizationConfig):
                return XpuMxfp4MoEMethod(layer.moe_config)
            else:
                quant_method = Mxfp4MoEMethod(layer.moe_config)
-                quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
                return quant_method
        elif isinstance(layer, Attention):
            # TODO: Add support for MXFP4 Attention.
@@ -243,7 +243,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
        self.weight_dtype = "mxfp4"
        self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)

-        self.marlin_input_dtype = None
        self.max_capture_size = (
            get_current_vllm_config().compilation_config.max_cudagraph_capture_size
        )
@@ -254,6 +253,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
            "Please check your environment and try again."
        )
        self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
+        self.moe_mk: mk.FusedMoEModularKernel | None = None

    def create_weights(
        self,
@@ -408,7 +408,30 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

    def process_weights_after_loading(self, layer):
        if self.mxfp4_backend == Mxfp4Backend.MARLIN:
-            prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)
+            prepare_moe_fp4_layer_for_marlin(
+                layer, input_dtype=get_marlin_input_dtype()
+            )
+
+            self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+            assert self.moe_quant_config is not None
+
+            prepare_finalize = maybe_make_prepare_finalize(
+                moe=self.moe,
+                quant_config=self.moe_quant_config,
+                routing_tables=layer._maybe_init_expert_routing_tables(),
+                allow_new_interface=True,
+            )
+            assert prepare_finalize is not None
+
+            self.moe_mk = mk.FusedMoEModularKernel(
+                prepare_finalize,
+                MarlinExperts(
+                    self.moe,
+                    self.moe_quant_config,
+                ),
+                inplace=not self.moe.disable_inplace,
+                shared_experts=None,
+            )
        elif (
            self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
            or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
@@ -910,27 +933,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
            raise NotImplementedError("EPLB is not supported for mxfp4")

        if self.mxfp4_backend == Mxfp4Backend.MARLIN:
-            return fused_marlin_moe(
-                x,
-                layer.w13_weight,
-                layer.w2_weight,
-                layer.w13_bias,
-                layer.w2_bias,
-                layer.w13_weight_scale,
-                layer.w2_weight_scale,
-                topk_weights,
-                topk_ids,
-                global_scale1=None,
-                global_scale2=None,
-                quant_type_id=scalar_types.float4_e2m1f.id,
-                apply_router_weight_on_input=layer.apply_router_weight_on_input,
-                global_num_experts=layer.global_num_experts,
-                activation=layer.activation,
-                expert_map=layer.expert_map,
-                input_dtype=self.marlin_input_dtype,
-                inplace=not self.moe.disable_inplace,
-            )
+            assert self.moe_mk is not None

+            return self.moe_mk(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=layer.activation,
+                global_num_experts=layer.global_num_experts,
+                expert_map=layer.expert_map,
+                apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            )
        assert _can_support_mxfp4(
            layer.use_grouped_topk,
            layer.topk_group,
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from typing import Any

 import numpy as np
@@ -11,7 +11,8 @@ from tqdm import tqdm
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
-from vllm.forward_context import set_forward_context
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import AttentionMetadataBuilder
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import (
@@ -34,14 +35,27 @@ class CudaGraphManager:
        self.max_num_reqs = self.scheduler_config.max_num_seqs
        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
        self.dp_size = vllm_config.parallel_config.data_parallel_size
+
+        self.uniform_decode_query_len = 1
+        spec_config = vllm_config.speculative_config
+        if spec_config is not None:
+            self.uniform_decode_query_len += spec_config.num_speculative_tokens
+
        self.compilation_config = vllm_config.compilation_config
        assert self.compilation_config is not None
        self.cudagraph_mode = self.compilation_config.cudagraph_mode
-        self.cudagraph_sizes = get_cudagraph_sizes(
+
+        use_uniform_decode_cudagraph = (
+            self.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
+            and self.cudagraph_mode.separate_routine()
+        )
+        self.cudagraph_sizes, self.uniform_decode_cudagraph_sizes = get_cudagraph_sizes(
            self.compilation_config.cudagraph_capture_sizes,
            self.max_num_reqs,
            self.max_num_tokens,
            self.cudagraph_mode,
+            self.uniform_decode_query_len,
+            use_uniform_decode_cudagraph,
        )

        self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
@@ -54,20 +68,16 @@ class CudaGraphManager:
        return len(self.cudagraph_sizes) > 0

    def get_cudagraph_size(
-        self,
-        num_tokens_after_padding: int,
-        num_tokens_per_request: Iterable[int],
+        self, num_tokens: int, uniform_decode: bool = False
    ) -> int | None:
-        return get_cudagraph_size(
-            num_tokens_after_padding,
-            num_tokens_per_request,
-            self.cudagraph_sizes,
-            self.cudagraph_mode,
-        )
+        if uniform_decode and self.uniform_decode_cudagraph_sizes:
+            return self.uniform_decode_cudagraph_sizes.get(num_tokens)
+        return self.cudagraph_sizes.get(num_tokens)

    def capture_graph(
        self,
        num_tokens: int,
+        capture_cg_mode: CUDAGraphMode,
        model: nn.Module,
        input_buffers: InputBuffers,
        mrope_positions: torch.Tensor | None,
@@ -75,8 +85,25 @@ class CudaGraphManager:
        block_tables: BlockTables,
        attn_metadata_builders: list[AttentionMetadataBuilder],
        kv_cache_config: KVCacheConfig,
+        has_lora: bool = False,
+        uniform_decode: bool = False,
    ) -> None:
-        num_reqs = min(num_tokens, self.max_num_reqs)
+        # select and check capture function
+        assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
+            f"Invalid capture_cudagraph_mode for capture: {capture_cg_mode}"
+        )
+        if capture_cg_mode == CUDAGraphMode.PIECEWISE:
+            capture_fn = self._capture_piecewise_graph
+        else:
+            capture_fn = self._capture_full_graph
+        # prepare inputs
+        if uniform_decode:
+            num_reqs = min(
+                cdiv(num_tokens, self.uniform_decode_query_len),
+                self.max_num_reqs,
+            )
+        else:
+            num_reqs = min(num_tokens, self.max_num_reqs)
        input_ids = input_buffers.input_ids[:num_tokens]
        positions = input_buffers.positions[:num_tokens]
        if self.uses_mrope:
@@ -92,6 +119,9 @@ class CudaGraphManager:
            attn_metadata_builders,
            self.max_model_len,
            kv_cache_config,
+            uniform_decode_query_len=(
+                self.uniform_decode_query_len if uniform_decode else 0
+            ),
        )
        num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)

@@ -112,13 +142,40 @@ class CudaGraphManager:
            if self.hidden_states is None:
                self.hidden_states = torch.empty_like(hidden_states)

+        capture_fn(
+            num_tokens=num_tokens,
+            num_reqs=num_reqs,
+            model=model,
+            input_ids=input_ids,
+            positions=positions,
+            inputs_embeds=inputs_embeds,
+            num_tokens_across_dp=num_tokens_across_dp,
+            attn_metadata=attn_metadata,
+            slot_mappings=slot_mappings,
+            has_lora=has_lora,
+        )
+
+    def _capture_full_graph(
+        self,
+        num_tokens: int,
+        num_reqs: int,
+        model: nn.Module,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None,
+        num_tokens_across_dp: torch.Tensor,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        has_lora: bool = False,
+    ) -> None:
+        assert attn_metadata is not None
        # Capture the graph.
        assert num_tokens not in self.graphs
        graph = torch.cuda.CUDAGraph()
        with (
            set_forward_context(
-                attn_metadata,
-                self.vllm_config,
+                attn_metadata=attn_metadata,
+                vllm_config=self.vllm_config,
                num_tokens=num_tokens,
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
                num_tokens_across_dp=num_tokens_across_dp,
@@ -131,9 +188,42 @@ class CudaGraphManager:
                positions=positions,
                inputs_embeds=inputs_embeds,
            )
+            assert self.hidden_states is not None
            self.hidden_states[:num_tokens] = hidden_states
        self.graphs[num_tokens] = graph

+    def _capture_piecewise_graph(
+        self,
+        num_tokens: int,
+        num_reqs: int,
+        model: nn.Module,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None,
+        num_tokens_across_dp: torch.Tensor,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        has_lora: bool = False,
+    ) -> None:
+        # create batch descriptor for piecewise cudagraph dispatch key
+        batch_descriptor = BatchDescriptor(num_tokens=num_tokens, has_lora=has_lora)
+
+        # Capture run - CUDAGraphWrapper inside torch.compile will auto capture.
+        with set_forward_context(
+            attn_metadata=None,  # piecewise no need attn_metadata
+            vllm_config=self.vllm_config,
+            num_tokens=num_tokens,
+            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+            num_tokens_across_dp=num_tokens_across_dp,
+            batch_descriptor=batch_descriptor,
+            slot_mapping=slot_mappings,
+        ):
+            model(
+                input_ids=input_ids,
+                positions=positions,
+                inputs_embeds=inputs_embeds,
+            )
+
    @torch.inference_mode()
    def capture(
        self,
@@ -144,11 +234,11 @@ class CudaGraphManager:
        block_tables: BlockTables,
        attn_metadata_builders: list[AttentionMetadataBuilder],
        kv_cache_config: KVCacheConfig,
+        has_lora: bool = False,
    ) -> None:
-        capture_graphs(
-            self.cudagraph_sizes,
-            self.device,
-            self.capture_graph,
+        common_kwargs = dict(
+            device=self.device,
+            capture_fn=self.capture_graph,
            model=model,
            input_buffers=input_buffers,
            mrope_positions=mrope_positions,
@@ -156,10 +246,50 @@ class CudaGraphManager:
            block_tables=block_tables,
            attn_metadata_builders=attn_metadata_builders,
            kv_cache_config=kv_cache_config,
+            has_lora=has_lora,
        )

-    def run(self, num_tokens: int) -> torch.Tensor:
-        assert num_tokens in self.graphs
+        # Phase 1: Capture for mixed prefill-decode batches if needed.
+        mixed_mode = self.cudagraph_mode.mixed_mode()
+        if mixed_mode != CUDAGraphMode.NONE:
+            capture_graphs(
+                cudagraph_sizes=self.cudagraph_sizes,
+                capture_cudagraph_mode=mixed_mode,
+                desc=f"Capturing CUDA graphs (mixed, {mixed_mode.name})",
+                uniform_decode=False,
+                **common_kwargs,
+            )
+
+        # Phase 2: Capture FULL graphs for uniform decode batches if needed.
+        # This is only needed if we use a separate routine for decode batches
+        # and the decode_mode is FULL.
+        if self.uniform_decode_cudagraph_sizes:
+            capture_graphs(
+                cudagraph_sizes=self.uniform_decode_cudagraph_sizes,
+                capture_cudagraph_mode=CUDAGraphMode.FULL,
+                desc="Capturing CUDA graphs (decode, FULL)",
+                uniform_decode=True,
+                **common_kwargs,
+            )
+
+    def get_cudagraph_runtime_mode(
+        self, num_reqs: int, num_tokens: int, max_query_len: int
+    ) -> tuple[CUDAGraphMode, int | None]:
+        is_uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
+            num_tokens == max_query_len * num_reqs
+        )
+
+        cudagraph_size = self.get_cudagraph_size(num_tokens, is_uniform_decode)
+        if cudagraph_size is None:
+            cudagraph_mode = CUDAGraphMode.NONE
+        elif is_uniform_decode:
+            cudagraph_mode = self.cudagraph_mode.decode_mode()
+        else:
+            cudagraph_mode = self.cudagraph_mode.mixed_mode()
+        return cudagraph_mode, cudagraph_size
+
+    def run_fullgraph(self, num_tokens: int) -> torch.Tensor:
+        assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens"
        self.graphs[num_tokens].replay()
        assert self.hidden_states is not None
        return self.hidden_states[:num_tokens]
@@ -170,22 +300,18 @@ def get_cudagraph_sizes(
    max_num_reqs: int,
    max_num_tokens: int,
    cudagraph_mode: CUDAGraphMode,
-) -> dict[int, int]:
-    if not cudagraph_mode.has_full_cudagraphs():
-        return {}
+    uniform_decode_query_len: int = 1,
+    uniform_decode_cudagraph: bool = False,
+) -> tuple[dict[int, int], dict[int, int]]:
+    # Support both FULL and PIECEWISE cudagraph modes
+    if cudagraph_mode == CUDAGraphMode.NONE:
+        return {}, {}
    if not capture_sizes:
-        return {}
+        return {}, {}

    capture_sizes = sorted(capture_sizes)
-    # Limit the capture sizes to the max number of requests or tokens.
-    upper_bound = (
-        max_num_reqs
-        if cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY
-        else max_num_tokens
-    )
-    capture_sizes = [x for x in capture_sizes if x <= upper_bound]
    if not capture_sizes:
-        return {}
+        return {}, {}

    cudagraph_sizes: dict[int, int] = {}
    for i in range(1, capture_sizes[-1] + 1):
@@ -193,45 +319,34 @@ def get_cudagraph_sizes(
            if i <= x:
                cudagraph_sizes[i] = x
                break
-    return cudagraph_sizes

-
-def get_cudagraph_size(
-    num_tokens_after_dp_padding: int,
-    num_tokens_per_request: Iterable[int],
-    cudagraph_sizes: dict[int, int],
-    cudagraph_mode: CUDAGraphMode,
-) -> int | None:
-    if not cudagraph_mode.has_full_cudagraphs():
-        # No full CUDA graph is used.
-        return None
-
-    size = cudagraph_sizes.get(num_tokens_after_dp_padding)
-    if size is None:
-        # No CUDA graph for this size.
-        return None
-
-    is_mixed = any(x > 1 for x in num_tokens_per_request)
-    if is_mixed and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL:
-        # Prefill is included, and this mode doesn't use CUDA graph for it.
-        return None
-    return size
+    uniform_decode_cudagraph_sizes: dict[int, int] = {}
+    if uniform_decode_cudagraph:
+        max_num_tokens = max_num_reqs * uniform_decode_query_len
+        uniform_decode_cudagraph_sizes = {
+            k: v
+            for k, v in cudagraph_sizes.items()
+            if v <= max_num_tokens and v >= uniform_decode_query_len
+        }
+    return cudagraph_sizes, uniform_decode_cudagraph_sizes


 def capture_graphs(
    cudagraph_sizes: dict[int, int],
    device: torch.device,
    capture_fn: Callable,
+    capture_cudagraph_mode: CUDAGraphMode,
+    desc: str = "Capturing CUDA graphs",
    **capture_kwargs,
 ) -> None:
    # Capture larger graphs first.
    sizes_to_capture = sorted(set(cudagraph_sizes.values()), reverse=True)
    if is_global_first_rank():
-        sizes_to_capture = tqdm(sizes_to_capture, desc="Capturing CUDA graphs")
+        sizes_to_capture = tqdm(sizes_to_capture, desc=desc)

    with graph_capture(device=device):
        for size in sizes_to_capture:
-            capture_fn(size, **capture_kwargs)
+            capture_fn(size, capture_cudagraph_mode, **capture_kwargs)


 def prepare_inputs_to_capture(
@@ -242,8 +357,12 @@ def prepare_inputs_to_capture(
    attn_metadata_builders: list[AttentionMetadataBuilder],
    max_model_len: int,
    kv_cache_config: KVCacheConfig,
+    uniform_decode_query_len: int = 0,
 ) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
-    num_tokens_per_req = num_tokens // num_reqs
+    if uniform_decode_query_len > 0:
+        num_tokens_per_req = uniform_decode_query_len
+    else:
+        num_tokens_per_req = num_tokens // num_reqs

    query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) * num_tokens_per_req
    query_start_loc_np[-1] = num_tokens
@@ -13,48 +13,65 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N


 def get_batch_metadata_across_dp(
-    num_tokens: int, cudagraph_size: int, dp_size: int, dp_rank: int
-) -> tuple[torch.Tensor, torch.Tensor]:
+    num_tokens: int,
+    cudagraph_size: int,
+    cudagraph_runtime_mode: int,
+    dp_size: int,
+    dp_rank: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    assert dp_size > 1
    # Use CPU group to avoid CPU-GPU synchronization.
    group = get_dp_group().cpu_group
-    tensor = torch.zeros(2, dp_size, dtype=torch.int32, device="cpu")
+    tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu")
    tensor[0][dp_rank] = num_tokens
    tensor[1][dp_rank] = cudagraph_size
+    tensor[2][dp_rank] = cudagraph_runtime_mode
    dist.all_reduce(tensor, group=group)
-    return tensor[0], tensor[1]
+    return tensor[0], tensor[1], tensor[2]


 def get_cudagraph_and_dp_padding(
-    num_tokens: int, cudagraph_size: int | None, dp_size: int, dp_rank: int
-) -> tuple[bool, int, torch.Tensor | None]:
+    num_tokens: int,
+    cudagraph_size: int | None,
+    cudagraph_runtime_mode: int,
+    dp_size: int,
+    dp_rank: int,
+) -> tuple[int, torch.Tensor | None, int]:
    if dp_size == 1:
        if cudagraph_size is not None:
-            return True, cudagraph_size, None
+            return cudagraph_size, None, cudagraph_runtime_mode
        else:
-            return False, num_tokens, None
+            return num_tokens, None, cudagraph_runtime_mode

+    # Convert None to -1 for sync (indicates no cudagraph available)
    if num_tokens == 0:
        cudagraph_size = 0
    elif cudagraph_size is None:
        cudagraph_size = -1
-    num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(
-        num_tokens, cudagraph_size, dp_size, dp_rank
+
+    num_tokens_across_dp, cudagraph_size_across_dp, cudagraph_mode_across_dp = (
+        get_batch_metadata_across_dp(
+            num_tokens, cudagraph_size, cudagraph_runtime_mode, dp_size, dp_rank
+        )
    )
    if torch.all(num_tokens_across_dp == 0).item():
        # All ranks have zero tokens to run.
-        return False, 0, None
+        return 0, None, 0

-    if torch.all(cudagraph_size_across_dp != -1).item():
-        # All ranks use CUDA graph or have zero tokens.
-        # Use CUDA graph for all ranks.
-        # Pad all ranks to the maximum CUDA graph size.
+    # Synchronize cudagraph_runtime_mode across ranks by taking the minimum.
+    synced_cudagraph_mode = int(cudagraph_mode_across_dp.min().item())
+    # Check if all ranks have valid cudagraph_size.
+    all_have_cudagraph = torch.all(cudagraph_size_across_dp != -1).item()
+
+    if synced_cudagraph_mode != 0 and all_have_cudagraph:
+        # All ranks use cudagraph. Pad to max cudagraph_size.
        max_cudagraph_size = int(cudagraph_size_across_dp.max().item())
        num_tokens_across_dp[:] = max_cudagraph_size
-        return True, max_cudagraph_size, num_tokens_across_dp
+        return max_cudagraph_size, num_tokens_across_dp, synced_cudagraph_mode
    else:
-        # Some ranks do not use CUDA graph. Use eager mode for all ranks.
-        # No padding is needed except for ranks that have no tokens to run.
+        # Fall back to eager mode (no cudagraph).
+        # Either some rank doesn't have cudagraph size or mode is NONE.
+        synced_cudagraph_mode = 0
        num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
        num_tokens_after_padding = int(num_tokens_across_dp[dp_rank].item())
-        return False, num_tokens_after_padding, num_tokens_across_dp
+        return num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode
@@ -15,7 +15,7 @@ from vllm.distributed.parallel_state import (
    get_pp_group,
    prepare_communication_buffer_for_model,
 )
-from vllm.forward_context import set_forward_context
+from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import get_model_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -140,7 +140,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            self.do_spec_decode = False
            self.num_speculative_steps = 0
            self.speculator = None
-
        self.req_states = RequestState(
            max_num_reqs=self.max_num_reqs,
            max_model_len=self.max_model_len,
@@ -458,6 +457,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                block_tables=self.block_tables,
                attn_metadata_builders=self.attn_metadata_builders,
                kv_cache_config=self.kv_cache_config,
+                has_lora=self.lora_config is not None,
            )
            if self.do_spec_decode:
                self.speculator.capture_model()
@@ -884,19 +884,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return empty_output

-        # Get the CUDA graph size. None means no CUDA graph is used.
-        cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
-            scheduler_output.total_num_scheduled_tokens,
-            scheduler_output.num_scheduled_tokens.values(),
+        # Get local cudagraph mode and size.
+        local_cudagraph_mode, local_cudagraph_size = (
+            self.cudagraph_manager.get_cudagraph_runtime_mode(
+                num_reqs=len(scheduler_output.num_scheduled_tokens),
+                num_tokens=scheduler_output.total_num_scheduled_tokens,
+                max_query_len=max(scheduler_output.num_scheduled_tokens.values()),
+            )
        )
-        use_cudagraph, num_tokens_after_padding, num_tokens_across_dp = (
+
+        # DP sync: num_tokens + cudagraph_size + cudagraph_mode
+        num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode = (
            get_cudagraph_and_dp_padding(
                scheduler_output.total_num_scheduled_tokens,
-                cudagraph_size,
+                local_cudagraph_size,
+                local_cudagraph_mode.value,
                self.parallel_config.data_parallel_size,
                self.parallel_config.data_parallel_rank,
            )
        )
+        cudagraph_runtime_mode = CUDAGraphMode(synced_cudagraph_mode)
        if num_tokens_after_padding == 0:
            # All DP ranks have zero tokens to run.
            empty_output = self.kv_connector.no_forward(scheduler_output)
@@ -946,16 +953,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            # FIXME(woosuk): Fix warmup for LoRA.

        # Run model.
-        if use_cudagraph:
-            # Run CUDA graph.
+        if cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            # Use explicit cudagraph replay for FULL mode.
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
            self.kv_connector.pre_forward(scheduler_output)
-            hidden_states = self.cudagraph_manager.run(
+            hidden_states = self.cudagraph_manager.run_fullgraph(
                input_batch.num_tokens_after_padding
            )
        else:
-            # Run PyTorch model in eager mode.
+            # For piecewise and eager mode, just call model().
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
@@ -970,13 +977,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                inputs_embeds = None
                assert intermediate_tensors is not None

+            batch_descriptor = BatchDescriptor(
+                num_tokens=input_batch.num_tokens_after_padding,
+                has_lora=self.lora_config is not None,
+            )
+
            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
-                # TODO(woosuk): Support piecewise CUDA graph.
-                cudagraph_runtime_mode=CUDAGraphMode.NONE,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
                num_tokens_across_dp=num_tokens_across_dp,
+                batch_descriptor=batch_descriptor,
                slot_mapping=input_batch.slot_mappings,
            ):
                self.kv_connector.pre_forward(scheduler_output)
@@ -7,7 +7,7 @@ import torch.nn as nn

 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.forward_context import set_forward_context
+from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import get_model
 from vllm.triton_utils import tl, triton
@@ -103,14 +103,17 @@ class EagleSpeculator:
        attn_metadata: dict[str, Any] | None,
        slot_mappings: dict[str, torch.Tensor] | None,
        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
        with set_forward_context(
            attn_metadata,
            self.vllm_config,
            num_tokens=num_tokens,
-            cudagraph_runtime_mode=CUDAGraphMode.NONE,
+            cudagraph_runtime_mode=cudagraph_runtime_mode,
            num_tokens_across_dp=num_tokens_across_dp,
            slot_mapping=slot_mappings,
+            batch_descriptor=batch_descriptor,
        ):
            ret_hidden_states = self.model(
                input_ids=self.input_buffers.input_ids[:num_tokens],
@@ -127,9 +130,11 @@ class EagleSpeculator:
    def generate_draft(
        self,
        num_reqs: int,
+        num_tokens_padded: int,
        attn_metadata: dict[str, Any],
        slot_mappings: dict[str, torch.Tensor],
        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
    ) -> None:
        pos = self.input_buffers.positions[:num_reqs]
        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
@@ -137,8 +142,14 @@ class EagleSpeculator:
        for step in range(1, self.num_speculative_steps):
            # Run the eagle model.
            last_hidden_states, hidden_states = self.run_model(
-                num_reqs, attn_metadata, slot_mappings, num_tokens_across_dp
+                num_tokens_padded,
+                attn_metadata,
+                slot_mappings,
+                num_tokens_across_dp,
+                cudagraph_runtime_mode,
            )
+            last_hidden_states = last_hidden_states[:num_reqs]
+            hidden_states = hidden_states[:num_reqs]
            logits = self.model.compute_logits(last_hidden_states)

            # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
@@ -283,12 +294,14 @@ class EagleSpeculator:
        )

        cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
-        if cudagraph_size is not None:
-            # Run CUDA graph.
-            self.cudagraph_manager.run(cudagraph_size)
+        cudagraph_mode = self.cudagraph_manager.cudagraph_mode
+        if cudagraph_size is not None and cudagraph_mode == CUDAGraphMode.FULL:
+            # Run full CUDA graph.
+            self.cudagraph_manager.run_fullgraph(cudagraph_size)
            return self.draft_tokens[:num_reqs]

-        # Run eager mode.
+        # Run eager or piecewise CUDA graph.
+        num_tokens_padded = cudagraph_size if cudagraph_size is not None else num_reqs
        query_start_loc_cpu = torch.arange(
            num_reqs + 1, dtype=torch.int32, device="cpu"
        )
@@ -312,8 +325,13 @@ class EagleSpeculator:
            slot_mappings, self.kv_cache_config
        )
        self.generate_draft(
-            num_reqs, attn_metadata, slot_mappings_by_layer, num_tokens_across_dp=None
-        )  # FIXME
+            num_reqs,
+            num_tokens_padded,
+            attn_metadata,
+            slot_mappings_by_layer,
+            num_tokens_across_dp=None,  # FIXME
+            cudagraph_runtime_mode=cudagraph_mode,
+        )
        return self.draft_tokens[:num_reqs]


@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Callable
+from typing import Any

 import torch

@@ -31,16 +32,17 @@ class EagleCudaGraphManager:
        self.compilation_config = vllm_config.compilation_config
        assert self.compilation_config is not None

-        self.cudagraph_mode = self.compilation_config.cudagraph_mode
-        if self.cudagraph_mode == CUDAGraphMode.FULL:
-            # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.
-            self.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
+        # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.
+        self.cudagraph_mode = self.compilation_config.cudagraph_mode.decode_mode()

-        self.cudagraph_sizes = get_cudagraph_sizes(
+        # only need to capture uniform decode cudagraph sizes (the 2nd return value)
+        _, self.cudagraph_sizes = get_cudagraph_sizes(
            self.compilation_config.cudagraph_capture_sizes,
            self.max_num_reqs,
            self.max_num_tokens,
            self.cudagraph_mode,
+            uniform_decode_query_len=1,
+            uniform_decode_cudagraph=True,
        )

        self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
@@ -54,12 +56,21 @@ class EagleCudaGraphManager:
    def capture_graph(
        self,
        num_tokens: int,
+        capture_cg_mode: CUDAGraphMode,
        generate_fn: Callable,
        input_buffers: InputBuffers,
        block_tables: BlockTables,
        attn_metadata_builders: list[AttentionMetadataBuilder],
        kv_cache_config: KVCacheConfig,
    ) -> None:
+        assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
+            f"Invalid capture_cudagraph_mode for capture: {capture_cg_mode}"
+        )
+        if capture_cg_mode == CUDAGraphMode.PIECEWISE:
+            capture_fn = self._capture_piecewise_graph
+        else:
+            capture_fn = self._capture_full_graph
+
        num_reqs = min(num_tokens, self.max_num_reqs)
        attn_metadata, slot_mappings = prepare_inputs_to_capture(
            num_reqs,
@@ -69,19 +80,70 @@ class EagleCudaGraphManager:
            attn_metadata_builders,
            self.max_model_len,
            kv_cache_config,
+            uniform_decode_query_len=1,
        )
        num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)

        # Warm up.
-        generate_fn(num_tokens, attn_metadata, slot_mappings, num_tokens_across_dp)
+        generate_fn(
+            num_reqs,
+            num_tokens,
+            attn_metadata,
+            slot_mappings,
+            num_tokens_across_dp,
+            CUDAGraphMode.NONE,
+        )

        # Capture the graph.
+        capture_fn(
+            num_reqs=num_reqs,
+            num_tokens=num_tokens,
+            generate_fn=generate_fn,
+            attn_metadata=attn_metadata,
+            slot_mappings=slot_mappings,
+            num_tokens_across_dp=num_tokens_across_dp,
+        )
+
+    def _capture_full_graph(
+        self,
+        num_reqs: int,
+        num_tokens: int,
+        generate_fn: Callable,
+        attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
+        num_tokens_across_dp: torch.Tensor,
+    ) -> None:
        assert num_tokens not in self.graphs
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph, self.pool):
-            generate_fn(num_tokens, attn_metadata, slot_mappings, num_tokens_across_dp)
+            generate_fn(
+                num_reqs,
+                num_tokens,
+                attn_metadata,
+                slot_mappings,
+                num_tokens_across_dp,
+                CUDAGraphMode.NONE,
+            )
        self.graphs[num_tokens] = graph

+    def _capture_piecewise_graph(
+        self,
+        num_reqs: int,
+        num_tokens: int,
+        generate_fn: Callable,
+        attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
+        num_tokens_across_dp: torch.Tensor,
+    ) -> None:
+        generate_fn(
+            num_reqs,
+            num_tokens,
+            attn_metadata,
+            slot_mappings,
+            num_tokens_across_dp,
+            CUDAGraphMode.PIECEWISE,
+        )
+
    @torch.inference_mode()
    def capture(
        self,
@@ -91,10 +153,15 @@ class EagleCudaGraphManager:
        attn_metadata_builders: list[AttentionMetadataBuilder],
        kv_cache_config: KVCacheConfig,
    ) -> None:
+        if self.cudagraph_mode == CUDAGraphMode.NONE:
+            return
+
        capture_graphs(
            self.cudagraph_sizes,
            self.device,
            self.capture_graph,
+            capture_cudagraph_mode=self.cudagraph_mode,
+            desc=f"Capturing eagle CUDA graphs ({self.cudagraph_mode.name})",
            generate_fn=generate_fn,
            input_buffers=input_buffers,
            block_tables=block_tables,
@@ -102,6 +169,6 @@ class EagleCudaGraphManager:
            kv_cache_config=kv_cache_config,
        )

-    def run(self, num_tokens: int) -> None:
+    def run_fullgraph(self, num_tokens: int) -> None:
        assert num_tokens in self.graphs
        self.graphs[num_tokens].replay()
Author	SHA1	Message	Date
Lucas Wilkinson	07782ff2fc	fix activation Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2026-02-19 11:09:03 -05:00
Lucas Wilkinson	0295b4555c	remove wheel build dependency on extensions Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2026-02-18 23:57:57 -05:00
Lucas Wilkinson	702e82e422	add 12.8 wheel Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2026-02-18 22:54:18 -05:00
Woosuk KwonandGitHub	c878b43b64	[Model Runner V2] Remove unnecessary copies in PW CUDA graph capture (#34849 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-02-18 15:52:50 -08:00
rasmithandGitHub	2b84ac669c	[CI][AMD][BugFix] Use torch.testing.assert_close instead of assert torch.allclose in test_rocm_skinny_gemms.py (#34181 ) Signed-off-by: Randall Smith <Randall.Smith@amd.com>	2026-02-18 23:10:19 +00:00
zhrrrandGitHub	11d3976b88	[Model Runner V2] support piecewise & mixed cudagraph (#32771 ) Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>	2026-02-18 15:03:17 -08:00
Yongye Zhu GitHub Robert Shaw	40da9625a1	[MoE Refactor] Convert mxfp4 marlin into modular kernel format (#34588 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-18 14:37:14 -08:00
Flora Feng GitHub Flora Feng Nicolò Lucchesi	8d9babd4de	Fix empty tool_call_id in Anthropic messages API tool result conversion (#34745 ) Signed-off-by: <> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <sfeng33@h100-01.nemg-001.lab.rdu2.dc.redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>	2026-02-18 14:31:59 -08:00
Aaron HaoandGitHub	e99ba957ec	[BUG] Fixing Weight Sync unit test (#34841 ) Signed-off-by: ahao-anyscale <ahao@anyscale.com>	2026-02-18 17:20:10 -05:00
Kyle SayersandGitHub	64ac1395e8	[Docs] Clean up speculators docs (#34065 ) Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>	2026-02-18 13:48:11 -08:00
Cyrus Leung GitHub Michael Goin	61cf087680	[Bugfix] Fix lora tests (#34834 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-02-18 13:22:31 -08:00
Wenlong Wang GitHub Robert Shaw Robert Shaw	847a57cd12	[Bugfix][MoE Kernel] Fix incorrect routing selection for models without expert groups (e.g., MiniMax-M2.1) (#34673 ) Signed-off-by: wwl2755 <wangwenlong2755@gmail.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-18 13:03:24 -08:00