A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 14:45 UTC
[Bugfix] Fix Dtypes for Pynccl Wrapper (#33030)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Jan 26, 2026 at 15:09 UTC
43a013c3a29194f7b88b1b611b3b0067592b8c67
2 files changed
+8
-2
vllm/distributed/device_communicators/pynccl_wrapper.py
+6
-2
@@ -72,7 +72,8 @@ class ncclDataTypeEnum:
72
ncclFloat64 = 8
73
ncclDouble = 8
74
ncclBfloat16 = 9
75
- ncclNumTypes = 10
75
+ ncclFloat8e4m3 = 10
76
+ ncclNumTypes = 11
77
78
@classmethod
79
def from_torch(cls, dtype: torch.dtype) -> int:
@@ -92,9 +93,12 @@ class ncclDataTypeEnum:
93
return cls.ncclFloat64
94
if dtype == torch.bfloat16:
95
return cls.ncclBfloat16
96
+ if dtype == torch.float8_e4m3fn:
97
+ return cls.ncclFloat8e4m3
98
raise ValueError(
99
f"Unsupported dtype {dtype}: should be one of "
97
- f"int8, uint8, int32, int64, float16, float32, float64, bfloat16."
100
+ f"int8, uint8, int32, int64, float16, float32, float64, bfloat16,"
101
+ " float8e4m3."
102
)
103
104
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+2
@@ -224,6 +224,8 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin
224
a1q_scale = None
225
226
if is_nvfp4 and a1q_scale is not None:
227
+ if a1q_scale.element_size() == 1:
228
+ a1q_scale = a1q_scale.view(torch.uint8)
229
a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
230
231
return a1q, a1q_scale, None, topk_ids, topk_weights