transformers
2e979850 - [MoE] Fix Blackwell GPU crash with torch._grouped_mm on torch <= 2.8 (#48014)

Commit
29 days ago
[MoE] Fix Blackwell GPU crash with torch._grouped_mm on torch <= 2.8 (#48014) * [MoE] Fix Blackwell GPU crash with torch._grouped_mm on torch <= 2.8 (#47999) On PyTorch <= 2.8, torch._grouped_mm is only supported on Hopper (compute capability 9.x / SM90). On newer architectures like Blackwell (SM100 B200 / SM120 RTX 5090), torch.cuda.get_device_capability(...) >= (9, 0) evaluated to True on torch <= 2.8, causing PyTorch to raise RuntimeError: torch._grouped_mm is only supported on CUDA devices with compute capability = 9.0. This PR scopes the torch._grouped_mm capability check on torch < 2.9 to get_device_capability(...)[0] == 9 so that Blackwell devices fall back cleanly to transformers::grouped_mm_fallback on older PyTorch versions. Fixes #47999. * Delete tests/utils/test_moe.py --------- Co-authored-by: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>
Author
Parents
Loading