Skip to content

[Model Runner V2] FP32 gumbel sampling.#41775

Merged
WoosukKwon merged 17 commits into
vllm-project:mainfrom
PatchouliTIS:patchy/fp32_gumbel_pr
May 15, 2026
Merged

[Model Runner V2] FP32 gumbel sampling.#41775
WoosukKwon merged 17 commits into
vllm-project:mainfrom
PatchouliTIS:patchy/fp32_gumbel_pr

Conversation

@PatchouliTIS

@PatchouliTIS PatchouliTIS commented May 6, 2026

Copy link
Copy Markdown
Contributor

Purpose

Profile results on H20

Optimized Gumbel sampling precision behavior
In:
vllm/v1/worker/gpu/sample/gumbel.py
vllm/envs.py

This PR makes FP64 use in Gumbel sampling optional rather than mandatory.

FP32 becomes the default fast option unless FP64 is explicitly enabled, and the code avoids instability by making sure the random uniform is never allowed to hit zero before applying -log(-log(u)).

Nsys profile:

FP64 sampler
Clipboard_Screenshot_1777022230

FP32 sampler
Clipboard_Screenshot_1777022253

Test Plan

Test Result


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
@PatchouliTIS PatchouliTIS marked this pull request as ready for review May 6, 2026 02:35

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added the v1 label May 6, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces an optimization for the Gumbel sampler by allowing it to operate in FP32 instead of FP64, controlled by a new environment variable VLLM_SAMPLER_FP64_GUMBEL. While the high-level functions and the gumbel_block_argmax helper were updated, the _gumbel_sample_kernel signature was not modified to accept the new USE_FP64 parameter, which will lead to a TypeError at runtime.

Comment thread vllm/v1/worker/gpu/sample/gumbel.py
PatchouliTaisa added 2 commits May 6, 2026 16:34
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
@njhill njhill added the verified Run pre-commit for new contributors without triggering other tests label May 12, 2026

@njhill njhill left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @PatchouliTIS, this looks nice to me.

@TheEpicDolphin TheEpicDolphin left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changes look good to me too! I think using fp32 by default here is definitely the right move. Can you please double check that draft acceptance rates don't regress from using fp32 vs fp64?

Comment thread vllm/v1/worker/gpu/sample/gumbel.py Outdated
@WoosukKwon WoosukKwon added the ready ONLY add when PR is ready to merge/full CI is needed label May 12, 2026
@WoosukKwon

Copy link
Copy Markdown
Collaborator

Given this is not a temporary flag, what about having an engine flag (--use-fp64-gumbel or something like that) instead of the env variable?

PatchouliTaisa added 2 commits May 13, 2026 15:16
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
@mergify

mergify Bot commented May 13, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label May 13, 2026
@mergify mergify Bot removed the needs-rebase label May 13, 2026
@PatchouliTIS

Copy link
Copy Markdown
Contributor Author

Given this is not a temporary flag, what about having an engine flag (--use-fp64-gumbel or something like that) instead of the env variable?

Move the env variable into an engine flag.

@mergify

mergify Bot commented May 13, 2026

Copy link
Copy Markdown
Contributor

Hi @PatchouliTIS, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?
mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:
# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

PatchouliTaisa and others added 6 commits May 13, 2026 16:15
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Comment thread =4.5.1 Outdated
Comment on lines +1 to +9
Requirement already satisfied: pre-commit in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (4.5.1)
Requirement already satisfied: cfgv>=2.0.0 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from pre-commit) (3.5.0)
Requirement already satisfied: identify>=1.0.0 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from pre-commit) (2.6.16)
Requirement already satisfied: nodeenv>=0.11.1 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from pre-commit) (1.10.0)
Requirement already satisfied: pyyaml>=5.1 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from pre-commit) (6.0.3)
Requirement already satisfied: virtualenv>=20.10.0 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from pre-commit) (20.36.1)
Requirement already satisfied: distlib<1,>=0.3.7 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from virtualenv>=20.10.0->pre-commit) (0.4.0)
Requirement already satisfied: filelock<4,>=3.20.1 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from virtualenv>=20.10.0->pre-commit) (3.20.3)
Requirement already satisfied: platformdirs<5,>=3.9.1 in /data/home/patchychen/miniconda3/lib/python3.13/site-packages (from virtualenv>=20.10.0->pre-commit) (4.5.0)

@TheEpicDolphin TheEpicDolphin May 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think this file was added by accident.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

irrelevant file removed.

@WoosukKwon WoosukKwon left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@PatchouliTIS LGTM. Can you please remove the redundant file so that we can merge?

PatchouliTaisa added 2 commits May 15, 2026 09:38
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
@PatchouliTIS

Copy link
Copy Markdown
Contributor Author

@PatchouliTIS LGTM. Can you please remove the redundant file so that we can merge?

Done, irrelevant files removed.

@WoosukKwon WoosukKwon merged commit 0162596 into vllm-project:main May 15, 2026
79 checks passed
omerpaz95 pushed a commit to omerpaz95/vllm that referenced this pull request May 18, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
omerpaz95 pushed a commit to omerpaz95/vllm that referenced this pull request May 18, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
mfylcek pushed a commit to mfylcek/vllm that referenced this pull request May 19, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
@njhill njhill added the v2 label May 20, 2026
h1t35h pushed a commit to h1t35h/vllm that referenced this pull request May 21, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
wangxiyuan pushed a commit to vllm-project/vllm-ascend that referenced this pull request May 25, 2026
### What this PR does / why we need it?

This PR updates vllm-ascend main2main validation to:

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1

Main upstream changes and vllm-ascend adaptations:

1. vLLM PR: vllm-project/vllm#41775  
   `[Model Runner V2] FP32 gumbel sampling`

   Upstream changes:
   - Adds `use_fp64_gumbel` config / argument.
- Changes Gumbel sampling and rejection sampling paths to accept
`use_fp64`.
   - Makes FP32 Gumbel the default path and keeps FP64 optional.

   vllm-ascend adaptation:
- Update `vllm_ascend/worker/v2/sample/gumbel.py` to accept `use_fp64`.
- Update `vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py`
to accept `use_fp64`.
- Raise `NotImplementedError` for `use_fp64=True` on NPU because the
current NPU Triton path does not support FP64 Gumbel / rejection
sampling.

2. vLLM PR: vllm-project/vllm#41162  
   `[Model Runner V2] Rebuild attn metadata between draft decode steps`

   Upstream changes:
- Adds `output_processed_logits` / `output_processed_logits_col`
protocol for EAGLE draft sampling.
- Uses `output_processed_logits_col` to select the current speculative
draft step when writing draft logits.

   vllm-ascend adaptation:
- Add `output_processed_logits` and `output_processed_logits_col`
support in NPU `gumbel_sample`.
- Store processed logits before adding Gumbel noise so rejection
sampling can consume the draft logits.


3. vLLM PR: vllm-project/vllm#42692  
   `[Bugfix] DFlash FP8 KV-Cache`

   Upstream changes:
- Changes `SpecDecodeBaseProposer.token_arange_np` from default NumPy
integer dtype to `np.int32`.

   vllm-ascend adaptation:
- Update `vllm_ascend/spec_decode/eagle_proposer.py` to use
`np.arange(..., dtype=np.int32)`.
- Keep the Ascend-specific `max_num_tokens + 1` behavior for
`query_start_loc_cpu[:batch_size + 1]`.

4. No direct upstream vLLM PR

   vllm-ascend adaptation:
   - Fix `OP_LOGE` / `OP_LOGW` format warnings in:
     - `csrc/moe/chunk_fwd_o/tiling_base/error_log.h`
     - `csrc/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h`

   Reason:
- The old macros passed `std::string` to `%s`, which fails when the
current build treats format warnings as errors.


### Does this PR introduce _any_ user-facing change?

No.

This PR is for main2main compatibility validation and internal NPU
sampling / speculative decoding adaptation.

One behavior to note: if users explicitly enable `--use-fp64-gumbel` on
NPU, vllm-ascend will raise `NotImplementedError`.

### How was this patch tested?

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1
- Validation focus:
  - NPU Gumbel sampling
  - EAGLE / MTP speculative decoding
  - rejection sampling
  - custom op build

---------

Signed-off-by: shenzhao <shenzhao9@huawei.com>
Co-authored-by: shenzhao <shenzhao9@huawei.com>
Liuweixiong0118 pushed a commit to Liuweixiong0118/vllm that referenced this pull request Jun 1, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
Signed-off-by: Liuweixiong0118 <lwx34158427@gmail.com>
yilunh998 pushed a commit to yilunh998/vllm-ascend that referenced this pull request Jun 2, 2026
### What this PR does / why we need it?

This PR updates vllm-ascend main2main validation to:

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1

Main upstream changes and vllm-ascend adaptations:

1. vLLM PR: vllm-project/vllm#41775
   `[Model Runner V2] FP32 gumbel sampling`

   Upstream changes:
   - Adds `use_fp64_gumbel` config / argument.
- Changes Gumbel sampling and rejection sampling paths to accept
`use_fp64`.
   - Makes FP32 Gumbel the default path and keeps FP64 optional.

   vllm-ascend adaptation:
- Update `vllm_ascend/worker/v2/sample/gumbel.py` to accept `use_fp64`.
- Update `vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py`
to accept `use_fp64`.
- Raise `NotImplementedError` for `use_fp64=True` on NPU because the
current NPU Triton path does not support FP64 Gumbel / rejection
sampling.

2. vLLM PR: vllm-project/vllm#41162
   `[Model Runner V2] Rebuild attn metadata between draft decode steps`

   Upstream changes:
- Adds `output_processed_logits` / `output_processed_logits_col`
protocol for EAGLE draft sampling.
- Uses `output_processed_logits_col` to select the current speculative
draft step when writing draft logits.

   vllm-ascend adaptation:
- Add `output_processed_logits` and `output_processed_logits_col`
support in NPU `gumbel_sample`.
- Store processed logits before adding Gumbel noise so rejection
sampling can consume the draft logits.

3. vLLM PR: vllm-project/vllm#42692
   `[Bugfix] DFlash FP8 KV-Cache`

   Upstream changes:
- Changes `SpecDecodeBaseProposer.token_arange_np` from default NumPy
integer dtype to `np.int32`.

   vllm-ascend adaptation:
- Update `vllm_ascend/spec_decode/eagle_proposer.py` to use
`np.arange(..., dtype=np.int32)`.
- Keep the Ascend-specific `max_num_tokens + 1` behavior for
`query_start_loc_cpu[:batch_size + 1]`.

4. No direct upstream vLLM PR

   vllm-ascend adaptation:
   - Fix `OP_LOGE` / `OP_LOGW` format warnings in:
     - `csrc/moe/chunk_fwd_o/tiling_base/error_log.h`
     - `csrc/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h`

   Reason:
- The old macros passed `std::string` to `%s`, which fails when the
current build treats format warnings as errors.

### Does this PR introduce _any_ user-facing change?

No.

This PR is for main2main compatibility validation and internal NPU
sampling / speculative decoding adaptation.

One behavior to note: if users explicitly enable `--use-fp64-gumbel` on
NPU, vllm-ascend will raise `NotImplementedError`.

### How was this patch tested?

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1
- Validation focus:
  - NPU Gumbel sampling
  - EAGLE / MTP speculative decoding
  - rejection sampling
  - custom op build

---------

Signed-off-by: shenzhao <shenzhao9@huawei.com>
Co-authored-by: shenzhao <shenzhao9@huawei.com>
Signed-off-by: yilunh <hanyilun1@huawei.com>
mvanhorn pushed a commit to mvanhorn/vllm that referenced this pull request Jun 4, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
andakai pushed a commit to andakai/vllm that referenced this pull request Jun 4, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
knight0528 pushed a commit to knight0528/vllm that referenced this pull request Jun 8, 2026
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
LostFox11 pushed a commit to LostFox11/vllm-ascend that referenced this pull request Jun 15, 2026
### What this PR does / why we need it?

This PR updates vllm-ascend main2main validation to:

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1

Main upstream changes and vllm-ascend adaptations:

1. vLLM PR: vllm-project/vllm#41775  
   `[Model Runner V2] FP32 gumbel sampling`

   Upstream changes:
   - Adds `use_fp64_gumbel` config / argument.
- Changes Gumbel sampling and rejection sampling paths to accept
`use_fp64`.
   - Makes FP32 Gumbel the default path and keeps FP64 optional.

   vllm-ascend adaptation:
- Update `vllm_ascend/worker/v2/sample/gumbel.py` to accept `use_fp64`.
- Update `vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py`
to accept `use_fp64`.
- Raise `NotImplementedError` for `use_fp64=True` on NPU because the
current NPU Triton path does not support FP64 Gumbel / rejection
sampling.

2. vLLM PR: vllm-project/vllm#41162  
   `[Model Runner V2] Rebuild attn metadata between draft decode steps`

   Upstream changes:
- Adds `output_processed_logits` / `output_processed_logits_col`
protocol for EAGLE draft sampling.
- Uses `output_processed_logits_col` to select the current speculative
draft step when writing draft logits.

   vllm-ascend adaptation:
- Add `output_processed_logits` and `output_processed_logits_col`
support in NPU `gumbel_sample`.
- Store processed logits before adding Gumbel noise so rejection
sampling can consume the draft logits.


3. vLLM PR: vllm-project/vllm#42692  
   `[Bugfix] DFlash FP8 KV-Cache`

   Upstream changes:
- Changes `SpecDecodeBaseProposer.token_arange_np` from default NumPy
integer dtype to `np.int32`.

   vllm-ascend adaptation:
- Update `vllm_ascend/spec_decode/eagle_proposer.py` to use
`np.arange(..., dtype=np.int32)`.
- Keep the Ascend-specific `max_num_tokens + 1` behavior for
`query_start_loc_cpu[:batch_size + 1]`.

4. No direct upstream vLLM PR

   vllm-ascend adaptation:
   - Fix `OP_LOGE` / `OP_LOGW` format warnings in:
     - `csrc/moe/chunk_fwd_o/tiling_base/error_log.h`
     - `csrc/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h`

   Reason:
- The old macros passed `std::string` to `%s`, which fails when the
current build treats format warnings as errors.


### Does this PR introduce _any_ user-facing change?

No.

This PR is for main2main compatibility validation and internal NPU
sampling / speculative decoding adaptation.

One behavior to note: if users explicitly enable `--use-fp64-gumbel` on
NPU, vllm-ascend will raise `NotImplementedError`.

### How was this patch tested?

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1
- Validation focus:
  - NPU Gumbel sampling
  - EAGLE / MTP speculative decoding
  - rejection sampling
  - custom op build

---------

Signed-off-by: shenzhao <shenzhao9@huawei.com>
Co-authored-by: shenzhao <shenzhao9@huawei.com>
LostFox11 pushed a commit to LostFox11/vllm-ascend that referenced this pull request Jun 15, 2026
### What this PR does / why we need it?

This PR updates vllm-ascend main2main validation to:

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1

Main upstream changes and vllm-ascend adaptations:

1. vLLM PR: vllm-project/vllm#41775  
   `[Model Runner V2] FP32 gumbel sampling`

   Upstream changes:
   - Adds `use_fp64_gumbel` config / argument.
- Changes Gumbel sampling and rejection sampling paths to accept
`use_fp64`.
   - Makes FP32 Gumbel the default path and keeps FP64 optional.

   vllm-ascend adaptation:
- Update `vllm_ascend/worker/v2/sample/gumbel.py` to accept `use_fp64`.
- Update `vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py`
to accept `use_fp64`.
- Raise `NotImplementedError` for `use_fp64=True` on NPU because the
current NPU Triton path does not support FP64 Gumbel / rejection
sampling.

2. vLLM PR: vllm-project/vllm#41162  
   `[Model Runner V2] Rebuild attn metadata between draft decode steps`

   Upstream changes:
- Adds `output_processed_logits` / `output_processed_logits_col`
protocol for EAGLE draft sampling.
- Uses `output_processed_logits_col` to select the current speculative
draft step when writing draft logits.

   vllm-ascend adaptation:
- Add `output_processed_logits` and `output_processed_logits_col`
support in NPU `gumbel_sample`.
- Store processed logits before adding Gumbel noise so rejection
sampling can consume the draft logits.


3. vLLM PR: vllm-project/vllm#42692  
   `[Bugfix] DFlash FP8 KV-Cache`

   Upstream changes:
- Changes `SpecDecodeBaseProposer.token_arange_np` from default NumPy
integer dtype to `np.int32`.

   vllm-ascend adaptation:
- Update `vllm_ascend/spec_decode/eagle_proposer.py` to use
`np.arange(..., dtype=np.int32)`.
- Keep the Ascend-specific `max_num_tokens + 1` behavior for
`query_start_loc_cpu[:batch_size + 1]`.

4. No direct upstream vLLM PR

   vllm-ascend adaptation:
   - Fix `OP_LOGE` / `OP_LOGW` format warnings in:
     - `csrc/moe/chunk_fwd_o/tiling_base/error_log.h`
     - `csrc/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h`

   Reason:
- The old macros passed `std::string` to `%s`, which fails when the
current build treats format warnings as errors.


### Does this PR introduce _any_ user-facing change?

No.

This PR is for main2main compatibility validation and internal NPU
sampling / speculative decoding adaptation.

One behavior to note: if users explicitly enable `--use-fp64-gumbel` on
NPU, vllm-ascend will raise `NotImplementedError`.

### How was this patch tested?

- vLLM version: `v0.20.2`
- vLLM main commit:
vllm-project/vllm@1ac10f1
- vLLM diff:
vllm-project/vllm@0d4d334...1ac10f1
- Validation focus:
  - NPU Gumbel sampling
  - EAGLE / MTP speculative decoding
  - rejection sampling
  - custom op build

---------

Signed-off-by: shenzhao <shenzhao9@huawei.com>
Co-authored-by: shenzhao <shenzhao9@huawei.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed v1 v2 verified Run pre-commit for new contributors without triggering other tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants