Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Bugfix] Fix FP8 torch._scaled_mm fallback for torch>2.5 with CUDA<12.4 #10095

Merged
merged 2 commits into from
Nov 7, 2024
Merged

[Bugfix] Fix FP8 torch._scaled_mm fallback for torch>2.5 with CUDA<12.4 #10095

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bd3a238
Fix FP8 torch._scaled_mm fallback for torch>2.5
mgoin Nov 6, 2024
ec074f7
Remove cuda
mgoin Nov 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions vllm/model_executor/layers/quantization/utils/w8a8_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,7 @@

# Input scaling factors are no longer optional in _scaled_mm starting
# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
TORCH_DEVICE_IDENTITY = torch.ones(1).cuda() \
if current_platform.is_rocm() else None
Comment on lines -10 to -11
Copy link
Collaborator

Choose a reason for hiding this comment

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

Do you know why it was checking for is_rocm here?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

The assumption being made was is_rocm == torch 2.5. If we really wanted to we could check the torch version, but I think we can just assume torch 2.5 as a floor now

TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)


def cutlass_fp8_supported() -> bool:
Expand Down Expand Up @@ -166,8 +165,7 @@ def apply_fp8_linear(

# Making sure the dummy tensor is on the same device as the weight
global TORCH_DEVICE_IDENTITY
if (TORCH_DEVICE_IDENTITY is not None
and TORCH_DEVICE_IDENTITY.device != weight.device):
if TORCH_DEVICE_IDENTITY.device != weight.device:
TORCH_DEVICE_IDENTITY = TORCH_DEVICE_IDENTITY.to(weight.device)

# GEMM
Expand Down