Support FP8 NVLS multimem reduction#826
Conversation
multimem.ld_reduce on FP8 inputs can use the .acc::f16 instruction variant for FP16 accumulation. Thread AccumT through the NVLS zero-copy allreduce path so callers can select that behavior via the existing accum dtype dispatch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Probe whether the current CUDA target supports FP8 multimem reduction and use that result when selecting or launching NVLS FP8 algorithms. Keep the existing algorithm sizing and buffer behavior unchanged for this sub-PR. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Apply formatter output and mark the unused extras parameter in the NVLS block-pipeline guard path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
This PR extends NVLS-based allreduce to better support FP8 by adding an FP16 accumulation path for multimem load+reduce, and by probing/gating FP8 NVLS algorithm selection based on device/runtime support (plumbed through the NCCL algorithm selector).
Changes:
- Add an FP8 NVLS support probe (
isFp8NvlsSupported) and propagate anfp8NvlsSupportedflag through the NCCL algorithm selector config. - Add an optional FP16 accumulator variant for FP8 multimem load+reduce (via
.acc::f16instruction variants). - Thread FP8 NVLS support gating into NVLS allreduce algorithm implementations.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/ext/nccl/nccl.cc | Computes FP8 NVLS support and passes it into the selector config. |
| src/ext/nccl/CMakeLists.txt | Adds collectives include path so NCCL extension can use collective utilities. |
| src/ext/nccl/algorithm_selector.hpp | Extends selector config with fp8NvlsSupported. |
| src/ext/nccl/algorithm_selector.cc | Gates NVLS selection for FP8 based on native FP8 and probed support. |
| src/ext/collectives/include/collective_utils.hpp | Declares FP8 dtype helpers and FP8 NVLS support probe API. |
| src/ext/collectives/include/allreduce/common.hpp | Extends multimem reduce helper to accept an accumulation type. |
| src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp | Adds cached fp8NvlsSupported_ state to the algorithm class. |
| src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp | Adds cached fp8NvlsSupported_ state to the algorithm class. |
| src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp | Adds cached fp8NvlsSupported_ state to the algorithm class. |
| src/ext/collectives/collective_utils.cu | Implements FP8 dtype checks and runtime FP8 NVLS support probing. |
| src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu | Adds AccumT plumbing for multimem load+reduce and runtime FP8 support gating. |
| src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu | Adds runtime FP8 NVLS support gating for this algorithm variant. |
| src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu | Adds runtime FP8 NVLS support gating for this algorithm variant. |
| include/mscclpp/switch_channel_device.hpp | Adds FP16-accumulator multimem load+reduce variants for FP8 vector types. |
There was a problem hiding this comment.
This PR adds FP16 accumulation for the multimem.ld_reduce path when the data type is FP8 (via the optional AccumT template parameter on SwitchChannelDeviceHandle::multimemLoadReduce and the dispatchFp8Accum helper), and threads an accumDtype parameter (default DataType::AUTO) through Algorithm::execute.
What this PR intentionally leaves for a future PR is wiring the accumulation type into algorithm selection. Today the accumulation type only influences the kernel dispatch after an algorithm has already been chosen; the selector itself does not know about it. Follow-up work should close that gap:
Carry accumDtype in the request. Add an accumDtype field to CollectiveRequest and thread it from the public/NCCL entry point (where it is currently defaulted to AUTO) down into selection, instead of only into execute().
Let algorithms declare accumulation-type support. Extend the Algorithm interface (e.g. a supportsAccumDtype(dtype, accumDtype) predicate, a capability set, or a new field in Constraint/tags) so each algorithm advertises which (dtype, accumDtype) combinations it can handle. For example, the NVLS/multimem path can support FP8 data with FP16 or FP32 accumulation, while packet/other paths may only support native accumulation.
Filter/prefer during selection. Update matchExecutionPlan and the native selectors (selectSingleNodeAllreduceBlackwell / any AlgoSelectFunc) to only pick algorithms that support the requested accumDtype, with a well-defined fallback (e.g. fall back to AUTO/native accumulation) when no candidate supports it.
Approving as-is since the accumulation plumbing is self-contained and correct; the selection changes above can be handled in a follow-up.
Summary
Validation