[Core/DBO][2/N] Dual-Batch Overlap add DeepEP High Throughput support and Prefill support #24845

LucasWilkinson · 2025-09-15T02:22:16Z

Purpose

Add support for having prefill requests inside ubatches; this is done by allowing a request to be split across ubatches with the second batch effectively chunked-prefill continuing on from the the first (since it always runs after)
Add support for the DeepEP high-throughput kernels with SM control for DeepGEMM and DeepEP (shout-out to @yewentao256 for the HT support)
General cleanups

Test Plan

lm_eval

Test Result

export VLLM_ALL2ALL_BACKEND=deepep_high_throughput

lm_eval --model local-completions --model_args "base_url=http://0.0.0.0:8000/v1/completions,model=deepseek-ai/DeepSeek-R1,num_concurrent=256" --tasks gsm8k
....
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9522|±  |0.0059|
|     |       |strict-match    |     5|exact_match|↑  |0.9500|±  |0.0060|

export VLLM_ALL2ALL_BACKEND=deepep_low_latency

lm_eval --model local-completions --model_args "base_url=http://0.0.0.0:8000/v1/completions,model=deepseek-ai/DeepSeek-R1,num_concurrent=256" --tasks gsm8k
....
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9522|±  |0.0059|
|     |       |strict-match    |     5|exact_match|↑  |0.9507|±  |0.0060|

HT Overlap Trace (2x8xH100)

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.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Signed-off-by: Sage Moore <sage@neuralmagic.com>

…inson/attn-slicing

Signed-off-by: Sage Moore <sage@neuralmagic.com>

…pping asserts Signed-off-by: Sage Moore <sage@neuralmagic.com>

…sult in an empty second ubatch Signed-off-by: Sage Moore <sage@neuralmagic.com>

…inson/attn-slicing

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Signed-off-by: Sage Moore <sage@neuralmagic.com>

…inson/attn-slicing

Signed-off-by: Sage Moore <sage@neuralmagic.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

yewentao256

Thanks for the work! A few more thoughts

yewentao256 · 2025-09-22T14:31:14Z

vllm/v1/worker/gpu_model_runner.py

+            allow_microbatching_options = [True, False] if \
+                capture_ubatched_graph else [False]


Could we simply use two bools?

yewentao256 · 2025-09-22T14:33:47Z

vllm/v1/worker/gpu_model_runner.py

        force_attention: bool = False,
        uniform_decode: bool = False,
-        allow_microbatching: bool = False,
+        allow_microbatching: bool = True,


In this case, would the param renamed to microbatching_fallback etc be better? (just feeling allow_microbatching doesn't show the idea you mention)
Or we can have detailed comments

yewentao256 · 2025-09-22T14:34:49Z

vllm/distributed/device_communicators/base_device_communicator.py

+        pass
+
+    def max_sms_used(self) -> Optional[int]:
+        return None  # None means it could use the whole GPU


Would -1 be better?

yewentao256 · 2025-09-22T14:35:22Z

vllm/engine/arg_utils.py

+        parallel_group.add_argument(
+            "--dbo-prefill-token-threshold",
+            **parallel_kwargs["dbo_prefill_token_threshold"])


Please add comments for the new arg

yewentao256 · 2025-09-22T14:38:47Z

vllm/model_executor/layers/fused_moe/modular_kernel.py

+            if hook is not None:
+                if dbo_enabled():
+                    # If DBO is being used, register the hook with the ubatch
+                    # context and call it in dbo_maybe_run_recv_hook instead of
+                    #  passing it to the receiver.
+                    dbo_register_recv_hook(hook)
+                    dbo_yield()
+                else:
+                    hook()
+
+            receiver()


Why we have this logic for two times?

Factor it out into a function, since it appears twice?

yewentao256 · 2025-09-22T14:38:58Z

vllm/model_executor/layers/fused_moe/modular_kernel.py

+                    dbo_yield()
+                else:
+                    hook()



Here once again

tlrmchlsmth

LGTM

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

tlrmchlsmth · 2025-09-22T14:57:01Z

vllm/model_executor/layers/fused_moe/modular_kernel.py

+            if hook is not None:
+                if dbo_enabled():
+                    # If DBO is being used, register the hook with the ubatch
+                    # context and call it in dbo_maybe_run_recv_hook instead of
+                    #  passing it to the receiver.
+                    dbo_register_recv_hook(hook)
+                    dbo_yield()
+                else:
+                    hook()
+
+            receiver()


Factor it out into a function, since it appears twice?

vllm/v1/worker/gpu_model_runner.py

…e.py Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

mergify · 2025-09-22T19:43:29Z

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

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

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

mergify · 2025-09-23T04:20:53Z

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

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

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

voipmonitor · 2025-09-24T02:09:51Z

@LucasWilkinson this commit introduces weired behaviour - the first http request with larger context is working normally but the subsequent requests are signifficantly slower. I have verified that it is this commit: cc1dc7e which should be this PR. a903669 is working normally

CUDA_VISIBLE_DEVICES=4,5 NCCL_DEBUG=INFO vllm serve /mnt/gpt-oss-120b --async-scheduling --tensor-parallel-size 2 --port 4999 --host 0.0.0.0 --config  GPT-OSS_Blackwell.yaml --served-model-name default --gpu-memory-utilization 0.8 --max-num-seqs 1024 --max-num-batched-tokens 32000

cat GPT-OSS_Blackwell.yaml
compilation-config: '{"pass_config":{"enable_fi_allreduce_fusion":false,"enable_noop":true},"custom_ops":["+rms_norm"],"cudagraph_mode":"FULL_AND_PIECEWISE"}'
async-scheduling: true
cuda-graph-sizes: 2048
max-num-batched-tokens: 32000

… and Prefill support (vllm-project#24845) Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

… and Prefill support (#24845) Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: yewentao256 <zhyanwentao@126.com>

varun-sundar-rabindranath · 2025-10-04T02:35:04Z

vllm/v1/worker/gpu_model_runner.py


-        if not should_ubatch:
-            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-        num_tokens += num_pad


removing this doesn't make the padding happen.

SageMoore and others added 30 commits June 2, 2025 19:04

format

92e0cc7

Signed-off-by: Sage Moore <sage@neuralmagic.com>

config format

44a595f

Signed-off-by: Sage Moore <sage@neuralmagic.com>

mla format

d4b502a

Signed-off-by: Sage Moore <sage@neuralmagic.com>

dp format

8332924

Signed-off-by: Sage Moore <sage@neuralmagic.com>

forward context format

243eac5

Signed-off-by: Sage Moore <sage@neuralmagic.com>

pplx format

d463976

Signed-off-by: Sage Moore <sage@neuralmagic.com>

fa format

e34e441

Signed-off-by: Sage Moore <sage@neuralmagic.com>

temporarily remove enable_microbatching

919eef9

Signed-off-by: Sage Moore <sage@neuralmagic.com>

temporarily remove enable_microbatching

2731e8c

Signed-off-by: Sage Moore <sage@neuralmagic.com>

Merge branch 'main' of https://github.com/neuralmagic/vllm into lwilk…

18e7d6c

…inson/attn-slicing

first round of fixes

539c0c3

Signed-off-by: Sage Moore <sage@neuralmagic.com>

more fixes

5f4a501

Signed-off-by: Sage Moore <sage@neuralmagic.com>

fix pplx a2a

e080e06

Signed-off-by: Sage Moore <sage@neuralmagic.com>

debugging

2e3484c

Signed-off-by: Sage Moore <sage@neuralmagic.com>

misc fixes. lm_eval still gets a wrong answer but it no longer hangs

f8848bb

Signed-off-by: Sage Moore <sage@neuralmagic.com>

added support for ubatch padding. not working

8a75b3a

Signed-off-by: Sage Moore <sage@neuralmagic.com>

ubatch padding should work now

a8675b7

Signed-off-by: Sage Moore <sage@neuralmagic.com>

more padding work. still gets the wrong answer

a00dabc

Signed-off-by: Sage Moore <sage@neuralmagic.com>

misc padding fixes

05ddc34

Signed-off-by: Sage Moore <sage@neuralmagic.com>

padding is getting correctness but there are still some edgecases tri…

60499f6

…pping asserts Signed-off-by: Sage Moore <sage@neuralmagic.com>

fix ubatch padding to account for the case where the padding would re…

e6e3407

…sult in an empty second ubatch Signed-off-by: Sage Moore <sage@neuralmagic.com>

Merge branch 'main' of https://github.com/neuralmagic/vllm into lwilk…

642bf2d

…inson/attn-slicing

fix using the same buffer across ubatches

ef3c01c

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

wip cudagraphs

d682f5e

more hacking

b74c731

Signed-off-by: Sage Moore <sage@neuralmagic.com>

misc changes

1d112d9

Signed-off-by: Sage Moore <sage@neuralmagic.com>

Merge branch 'main' of https://github.com/neuralmagic/vllm into lwilk…

0889f66

…inson/attn-slicing

more fixes

ff2dd13

Signed-off-by: Sage Moore <sage@neuralmagic.com>

setup deepepll for ubatching

a4def24

Signed-off-by: Sage Moore <sage@neuralmagic.com>

yields now work with deepep_ll

930efd0

Signed-off-by: Sage Moore <sage@neuralmagic.com>

LucasWilkinson added 3 commits September 20, 2025 19:58

Merge remote-tracking branch 'origin/main' into lwilkinson/dbo-prefill

68ebc6b

Merge remote-tracking branch 'origin/main' into lwilkinson/dbo-prefill

8dd1bb0

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

remove unrelated change

4781120

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

yewentao256 reviewed Sep 22, 2025

View reviewed changes

tlrmchlsmth approved these changes Sep 22, 2025

View reviewed changes

LucasWilkinson and others added 2 commits September 22, 2025 11:14

Update vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finaliz…

99f6c06

…e.py Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

review comments

9416c69

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

mergify bot added the needs-rebase label Sep 22, 2025

Merge branch 'main' into lwilkinson/dbo-prefill

4e9e2ce

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

mergify bot removed the needs-rebase label Sep 23, 2025

fix pre-commit

356ddcb

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

tlrmchlsmth enabled auto-merge (squash) September 23, 2025 02:37

mergify bot added the needs-rebase label Sep 23, 2025

Merge branch 'main' into lwilkinson/dbo-prefill

4c631db

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

mergify bot removed the needs-rebase label Sep 23, 2025

precommit fix

f0bc569

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

tlrmchlsmth mentioned this pull request Sep 23, 2025

[RFC]: Data Parallel Attention and Expert Parallel MoEs #16037

Closed

37 tasks

tlrmchlsmth merged commit cc1dc7e into vllm-project:main Sep 23, 2025
51 checks passed

jmkuebler mentioned this pull request Sep 24, 2025

[Bug] fix import and unit test #25558

Merged

5 tasks

Conless mentioned this pull request Sep 25, 2025

[Core] Nanoflow-style Computation-Communication Overlap #23592

Open

4 tasks

This was referenced Sep 30, 2025

Quick fix for IMA with the Prefix Prefill kernel during graph capture #25983

Merged

[Bug]: IMA with the prefix_prefill::context_attention_fwd kernel when capturing full cudagraphs #25985

Open

varun-sundar-rabindranath mentioned this pull request Oct 4, 2025

[BugFix] Pad input buffers in _dummy_run #26209

Draft

varun-sundar-rabindranath reviewed Oct 4, 2025

View reviewed changes

		allow_microbatching_options = [True, False] if \
		capture_ubatched_graph else [False]

Uh oh!

[Core/DBO][2/N] Dual-Batch Overlap add DeepEP High Throughput support and Prefill support #24845

[Core/DBO][2/N] Dual-Batch Overlap add DeepEP High Throughput support and Prefill support #24845

Uh oh!

Conversation

LucasWilkinson commented Sep 15, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

yewentao256 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tlrmchlsmth left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mergify bot commented Sep 22, 2025

Uh oh!

mergify bot commented Sep 23, 2025

Uh oh!

Uh oh!

voipmonitor commented Sep 24, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

LucasWilkinson commented Sep 15, 2025 •

edited by github-actions bot

Loading