vllm-project/vllm

[Bug]: AsyncTP pass has poor perf on B200

Open

#27893 opened on Oct 31, 2025

View on GitHub
 (5 comments) (2 reactions) (2 assignees)Python (80,034 stars) (16,816 forks)batch import
bughelp wantedkeep-opentorch.compile

Description

Your current environment

main

🐛 Describe the bug

When running a model with quant_fp8 + flashinfer on B200s, the matmul kernel that gets used is flashinfer_scaled_fp8_mm which turns into a vllm.bmm_fp8 op in the graph. This doesn't work with the AsyncTP pass in the torch.compile compilation as there doesn't exist a pattern/replacement for this.

Following the other patterns, I wrote one for bmm_fp8: https://github.com/vllm-project/vllm/pull/26933/commits/b0ab87b121acd1d4c52f3fbee12c3a447ea8f6b4

However, perf is a lot worse:

This might be because we're replacing the bmm_fp8 + reduce_scatter with torch.ops.symm_mem.patched_fused_scaled_matmul_reduce_scatter, but this op does not have a B200 specific implementation (it just calls into aten._scaled_mm)

cc @ProExpertProg @cascade812

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Contributor guide

Tech stack
pythonpytorch
Domain
machine learningperformance
Issue type
bug
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
5
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
over 1 week
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
blocked
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
mostly clear
Prerequisites
Understanding of vllm's distributed inferencePyTorch compiled mode experienceFamiliarity with GPU kernel fusionKnowledge of B200 architecture
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
10
Research direction
Investigate why the proposed pattern using `torch.ops.symm mem.patched fused scaled matmul reduce scatter` performs poorly on B200. This op calls `aten. scaled mm` which lacks B200 specific optimization. Compare with flashinfer's custom `bmm fp8` kernel. Look at existing async TP patterns in vllm (e.g., in `vllm/distributed/communication.py`) and study the commit referenced (b0ab87b) to understand the replacement logic. Consider whether a B200 specific fused kernel or alternative scheduling is needed.