[Bugfix] Fix Dtypes for Pynccl Wrapper

[robertgshaw2-redhat]

Purpose

• fixed [CI Failure]: MoE Integration Tests #33029
• recent PR [Perf][Kernel] Optimize FP4 quantization kernels (SM100F) #32520 swtiched scale type from uint8 to fp8
• this breaks FI AG/RS, because scales are now in a dtype that was not supported by the collectives
• also

Test Plan

• ci

---

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.

[gemini-code-assist]

Code Review

This pull request adds support for fp8 data types in the pynccl wrapper, which is necessary for distributed communication with fp8 tensors. The changes correctly add ncclFloat8e4m3 to the data type enum and handle torch.float8_e4m3fn in the type conversion logic.

My review includes a suggestion to also handle the torch.float8_e4m3fnuz variant, as it appears to be used in other parts of the codebase, to prevent potential runtime errors. This will make the fp8 support more robust.

[gemini-code-assist]

The codebase, for example in vllm/model_executor/layers/quantization/utils/fp8_utils.py, seems to use both torch.float8_e4m3fn and torch.float8_e4m3fnuz. This function should handle both types to avoid ValueError during collective communication operations with torch.float8_e4m3fnuz tensors. The error message is also updated for clarity.

For more complete FP8 support, you might also consider adding torch.float8_e5m2. This would involve adding ncclFloat8e5m2 to ncclDataTypeEnum and handling torch.float8_e5m2 in this method.

Suggested change

	if dtype == torch.float8_e4m3fn:
	return cls.ncclFloat8e4m3
	raise ValueError(
	f"Unsupported dtype {dtype}: should be one of "
	f"int8, uint8, int32, int64, float16, float32, float64, bfloat16."
	f"int8, uint8, int32, int64, float16, float32, float64, bfloat16,"
	" float8e4m3."
	)
	if dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz):
	return cls.ncclFloat8e4m3
	raise ValueError(
	f"Unsupported dtype {dtype}: should be one of "
	"int8, uint8, int32, int64, float16, float32, float64, bfloat16, "
	"float8_e4m3fn, "
	"float8_e4m3fnuz."
	)

[LucasWilkinson]

LGTM thanks for the quick fix!

[LopezCastroRoberto]

LGTM too. Thanks for the fix, Rob

[robertgshaw2-redhat]

updating main for CI