microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-14 19:27:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add configurable FP8 low-latency dispatch

Browse Source 
Quantize BF16 dispatch payloads to FP8 E4M3 with format-defined block scales while preserving BF16 expert outputs for combine. Clean up the sender structure, payload metadata, vector conversions, Python API, and multi-rank coverage.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
This commit is contained in:
Binyang Li
2026-07-13 18:47:21 +00:00
parent 7c1298dbae
commit 325f79f9dc
20 changed files with 847 additions and 487 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/mscclpp/ep/README.md
View File

@@ -261,10 +261,9 @@ can leave those fields as `None`.
```python
@dataclass
class QuantConfig:
dtype: Optional[torch.dtype] = None
format: Optional[DispatchDataType] = None
block_scales: Optional[torch.Tensor] = None
global_scale: Optional[torch.Tensor] = None
block_size: Optional[int] = None
class DispatchLayout(str, Enum):
@@ -473,17 +472,17 @@ combine to reduce the `K` expert results for each token back to `[T, H]`.
### `quant`
`quant` contains activation quantization metadata for `input`. It should be
`None` for BF16/FP16 input. The quantized tensor dtype is stored in
`quant.dtype`.
`None` for BF16/FP16 input. `quant.format` defines the tensor representation
and scale layout.
Examples:
| Format | `input` | `quant.dtype` | `quant.block_scales` | `quant.global_scale` |
|---|---|---|---|---|
| BF16/FP16 | `[T, H]` | `None` | `None` | `None` |
| FP8 E4M3 | `[T, H]` FP8 | `torch.float8_e4m3fn` | `[T, H / block_size]`, often block size 128 | usually `None` |
| NVFP4 | backend-defined packed/logical `[T, H]` | backend-defined | block scale tensor | optional global scale |
| MXFP8 | backend-defined `[T, H]` | backend-defined | micro-scale tensor, e.g. E8M0 blocks | optional/global if required |
| Format | `input` | `quant.block_scales` | `quant.global_scale` |
|---|---|---|---|
| BF16/FP16 | `[T, H]` | `None` | `None` |
| FP8 E4M3 | `[T, H]` FP8 | `[T, H / 128]` | usually `None` |
| NVFP4 | backend-defined packed/logical `[T, H]` | block scale tensor | optional global scale |
| MXFP8 | backend-defined `[T, H]` | micro-scale tensor, e.g. E8M0 blocks | optional/global if required |
The API should not assume quantization scale is a scalar. For FP8 paths in
DeepEP/SGLang, scales are usually per token and per hidden block.
@@ -504,7 +503,7 @@ output_buffer: [num_local_experts, world_size * max_tokens_per_rank, hidden]
The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16.
For FP8 dispatch it is FP8 and the returned `DispatchOutput.quant` carries the
matching dtype and scale tensor.
matching format and scale tensor.
`output_buffer` is required for LL because the MLP runner often owns or reuses
workspace memory. `MoECommunicator` writes dispatch output into the provided
@@ -816,15 +815,16 @@ output = moe_comm.combine(expert_output, handle)
Quantized path:
```python
moe_comm = MoECommunicator(
...,
quant=QuantConfig(format=DispatchDataType.FP8_E4M3),
)
recv, handle = moe_comm.dispatch(
input=x_fp8,
input=hidden_states, # BF16 input, quantized during dispatch
topk_ids=topk_ids,
weights=topk_weights,
quant=QuantConfig(
dtype=torch.float8_e4m3fn,
block_scales=x_scales,
block_size=128,
),
quant=None,
output_buffer=recv_buffer,
)

2
python/mscclpp/ep/__init__.py
View File

@@ -15,6 +15,7 @@ from .communicator import ( # noqa: F401
CombineContext,
CombineMode,
DispatchHandle,
DispatchDataType,
DispatchLayout,
DispatchLayoutInfo,
DispatchOutput,
@@ -38,6 +39,7 @@ __all__ = [
"CombineContext",
"CombineMode",
"DispatchHandle",
"DispatchDataType",
"DispatchLayout",
"DispatchLayoutInfo",
"DispatchOutput",

1
python/mscclpp/ep/_cpp.py
View File

@@ -15,6 +15,7 @@ except ImportError as exc: # pragma: no cover
DispatchLayout = _cpp.DispatchLayout
MoEMode = _cpp.MoEMode
CombineMode = _cpp.CombineMode
DispatchDataType = _cpp.DispatchDataType
Config = getattr(_cpp, "Config", None)

3
python/mscclpp/ep/communicator.py
View File

@@ -8,7 +8,7 @@ from typing import Optional, Tuple
import torch
from ._cpp import CombineMode, DispatchLayout, MoEMode
from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode
from .high_throughput import HighThroughputBackend
from .low_latency import LowLatencyBackend
from .types import (
@@ -36,6 +36,7 @@ __all__ = [
"CombineContext",
"CombineMode",
"DispatchHandle",
"DispatchDataType",
"DispatchLayout",
"DispatchLayoutInfo",
"DispatchOutput",

60
python/mscclpp/ep/low_latency.py
View File

@@ -8,7 +8,7 @@ from typing import Any, Optional
import torch
from ._cpp import CombineMode, DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint
from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint
from .types import (
DispatchHandle,
DispatchLayoutInfo,
@@ -22,6 +22,24 @@ from .types import (
from .utils import cuda_stream_ptr, resolve_expert_placement
def _resolve_dispatch_data_type(quant: Optional[QuantConfig]) -> DispatchDataType:
if quant is None:
return DispatchDataType.BF16
quant_format = quant.format
if quant_format is not None and not isinstance(quant_format, DispatchDataType):
raise TypeError("quant.format must be a DispatchDataType")
if quant_format is None:
raise ValueError("quant.format is required")
if quant_format == DispatchDataType.MXFP8_E4M3:
raise NotImplementedError("MXFP8 dispatch is reserved but not implemented")
if quant_format != DispatchDataType.FP8_E4M3:
raise ValueError("unsupported low-latency quantization format")
if quant.block_scales is not None or quant.global_scale is not None:
raise ValueError("communicator quant config must not contain precomputed scales")
return DispatchDataType.FP8_E4M3
class LowLatencyRuntime:
"""Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``)."""
@@ -109,12 +127,12 @@ class LowLatencyBackend:
if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank):
raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity")
if config.quant is not None:
raise NotImplementedError("low-latency quantization is not implemented yet")
self.dispatch_data_type = _resolve_dispatch_data_type(config.quant)
num_rdma_bytes = get_low_latency_rdma_size_hint(
self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts, self.topk
)
self._dispatch_scales: Optional[torch.Tensor] = None
self._dispatch_src_info: Optional[torch.Tensor] = None
self._dispatch_layout_range: Optional[torch.Tensor] = None
self._dispatch_count: Optional[torch.Tensor] = None
@@ -151,12 +169,13 @@ class LowLatencyBackend:
del previous_handle
self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer)
out_buf, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer)
out_buf, scales, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer)
self._runtime.cpp_runtime.dispatch(
input.data_ptr(),
topk_ids.data_ptr(),
0 if weights is None else weights.data_ptr(),
out_buf.data_ptr(),
0 if scales is None else scales.data_ptr(),
src_info.data_ptr(),
layout_range.data_ptr(),
count.data_ptr(),
@@ -165,12 +184,21 @@ class LowLatencyBackend:
self.topk,
self.max_tokens_per_rank,
self.num_experts,
self.dispatch_data_type,
self.num_blocks,
cuda_stream_ptr(stream),
)
output_quant = (
None
if scales is None
else QuantConfig(
format=self.dispatch_data_type,
block_scales=scales,
)
)
output_info = DispatchOutputInfo(
layout=DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count),
quant=None,
quant=output_quant,
)
dispatch_out = DispatchOutput(
tokens=out_buf,
@@ -218,6 +246,7 @@ class LowLatencyBackend:
self.topk,
context.num_max_dispatch_tokens_per_rank,
context.num_experts,
self.dispatch_data_type,
self.combine_mode,
self.num_blocks - 2,
cuda_stream_ptr(stream),
@@ -235,8 +264,16 @@ class LowLatencyBackend:
(self.num_local_experts, self.world_size), dtype=torch.int64, device=device
)
self._dispatch_count = torch.empty((self.num_local_experts,), dtype=torch.int32, device=device)
self._dispatch_scales = None
if self.dispatch_data_type == DispatchDataType.FP8_E4M3:
num_scales = self.hidden_size // 128
scale_storage = torch.empty(
(self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device
)
self._dispatch_scales = scale_storage.transpose(1, 2)
return (
output_buffer,
self._dispatch_scales,
self._dispatch_src_info,
self._dispatch_layout_range,
self._dispatch_count,
@@ -246,7 +283,9 @@ class LowLatencyBackend:
if output_buffer is None:
raise ValueError("output_buffer is required for low-latency dispatch")
if quant is not None:
raise NotImplementedError("low-latency quantization is not implemented yet")
raise NotImplementedError(
"per-call input quant metadata is not supported; configure dispatch output quantization on the communicator"
)
if input.dim() != 2 or not input.is_contiguous():
raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor")
if input.device.type != "cuda" or input.dtype != torch.bfloat16:
@@ -275,8 +314,9 @@ class LowLatencyBackend:
expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size)
if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous():
raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor")
if output_buffer.device != input.device or output_buffer.dtype != torch.bfloat16:
raise ValueError("output_buffer must be a BF16 CUDA tensor on the same device as input")
expected_dtype = torch.float8_e4m3fn if self.dispatch_data_type == DispatchDataType.FP8_E4M3 else torch.bfloat16
if output_buffer.device != input.device or output_buffer.dtype != expected_dtype:
raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input")
if tuple(output_buffer.shape) != expected_shape:
raise ValueError(f"output_buffer shape must be {expected_shape}")
@@ -286,6 +326,10 @@ class LowLatencyBackend:
context = handle.combine_context
if context.num_experts != self.num_experts or context.hidden_size != self.hidden_size:
raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration")
output_quant = handle.output_info.quant
handle_data_type = DispatchDataType.BF16 if output_quant is None else output_quant.format
if handle_data_type != self.dispatch_data_type:
raise ValueError("DispatchHandle quantization does not match this MoECommunicator configuration")
slots_per_expert = self.world_size * self.max_tokens_per_rank
if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR:
expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)

11
python/mscclpp/ep/types.py
View File

@@ -10,19 +10,22 @@ from typing import Any, List, Optional, Union
import torch
import mscclpp
from ._cpp import CombineMode, DispatchLayout, MoEMode
from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode
# Quantization metadata.
@dataclass
class QuantConfig:
"""Quantization metadata associated with an activation tensor."""
"""Quantization metadata associated with an activation tensor.
dtype: Optional[torch.dtype] = None
Low-latency FP8 dispatch returns ``block_scales`` with the activation's
leading dimensions and a format-defined final scale dimension.
"""
format: Optional[DispatchDataType] = None
block_scales: Optional[torch.Tensor] = None
global_scale: Optional[torch.Tensor] = None
block_size: Optional[int] = None
# Communicator construction.