[Feature] Support DeepSeek tile-wise (1×128) FP8 activation dispatch

[silencelamb]

Summary

Buffer.dispatch now accepts DeepEP-style prequantized FP8 activations as a (data, scales) tuple:

• data: contiguous [S, H] torch.float8_e4m3fn
• scales: [S, H // 128] torch.float32 dequantization scales, one per 1×128 tile (row-major or column-major); H must be divisible by 128

It returns (recv_data [NvS, H], recv_scales [NvS, H // 128]), both contiguous, in the same expert-grouped layout as BF16. Padding rows of both data and scales are zeroed. Combine still takes BF16 expert outputs.

(recv_data, recv_scales), route_weights_nvs, cu_seqlens, plan = buffer.dispatch(
    (data, scales), route_weights_sk, topk_experts_sk, tokens_per_expert, zero_copy=True,
)

Design

• No extra memory. The FP8 data and scale views live inside the existing BF16 communication allocation. Each rank's BF16 shard holds FP8 rows in its first half and scale rows in its second half, and the rank stride stays the same, so BF16 combine and backward dispatch keep working on the same buffer. zero_copy=True returns views for both tensors. The README explains that hidden_nvsh_buffer_view (the BF16 expert output) shares their storage.
• Dispatch kernel. Payload rows go through the existing TMA G2S/S2G pipeline as 1-byte elements. When the scales are contiguous, 16-B aligned and H % 512 == 0 (e.g. 3584, 7168), the producer also loads each token's scale row into a small smem ring with TMA, on the same mbarrier as the payload row. The consumer warp then scatters scales from smem with coalesced stores to the destinations that receive the payload (dst >= 0). Other layouts (column-major, misaligned, H % 512 != 0) fall back to per-element global loads.
• Dispatch epilogue. Duplicate expansion copies scale rows with the payload. Primary scale rows are staged by TMA in the same way.
• Unchanged contracts. Plan reuse, async_finish, router_weights_zero_copy, PDL and CUDA graphs all behave as before. All ranks must use the same activation format for a given call.
• Int64 row-index arithmetic in the dispatch kernel, because the FP8 view doubles the row count per rank stride.

The staged-scale path matters. With per-element loads on the consumer's critical path, scale handling took up to about one third of FP8 dispatch time (measured by disabling it). At H=3584 with skewed routing, FP8 was then only 0.99–1.16× faster than BF16.

Validation

8× H200 (NV18 full mesh), PyTorch 2.9.1+cu128, nvidia-cutlass-dsl 4.6.2, rebased on 33327eb (Public Release 26/09).

torchrun --standalone --nproc_per_node=8 -m pytest -q \
    tests/test_dispatch_fp8.py tests/test_dispatch.py tests/test_combine.py
torchrun --standalone --nproc_per_node=8 -m pytest -q tests/test_e2e.py

• 75 FP8 + existing BF16 dispatch/combine tests pass on every rank, and test_e2e.py passes. (test_e2e.py needs its own torchrun invocation. Collected together with test_dispatch.py, it fails with "initialize the default process group twice" on master as well.)
• The tests compare the bits of payload, scales and weights against a gathered reference, and check zeroed padding for both data and scales. They cover:
  • FP8 → BF16 → FP8 dispatch on the same plan
  • owned-output lifetime
  • async completion and PDL on/off
  • row-major (TMA path), column-major and misaligned (fallback path) scales
  • CUDA graphs on both scale paths
  • pipeline-stage reuse in both kernels
  • invalid inputs rejected before planning
  • FP8 dispatch → BF16 combine

Performance

benchmarks/bench_dispatch_fp8.py on 8× H200. E=256, K=8, token_padding=128, 32 dispatch SMs, PDL on. 48 configurations: S∈{256, 2048, 8192} × H∈{3584, 7168} × routing bias∈{0, 1} × fresh/cached plan × zero_copy∈{0, 1}. Inputs are prequantized, and timings include scales, duplicate expansion, planning (fresh) and output copies (zero_copy=0).

Metric (vs upstream master 33327eb)	Median	Range
Upstream BF16 latency / branch FP8 latency	1.55×	1.14–1.88×
Branch BF16 latency change vs upstream BF16	+0.08%	−0.71% to +4.09%*

*The +4.09% outlier (S=256, ~60 µs) is noise: three re-runs gave 59.1–60.6 µs upstream and 58.6–59.4 µs on the branch. The BF16 kernel code only changed in its index arithmetic.

FP8 is faster than BF16 in all 48 configurations. Selected rows (cached plan, zero_copy=1):

S	H	bias	upstream BF16 (µs)	FP8 (µs)	speedup
256	7168	0	115.7	68.5	1.69×
2048	3584	1	291.6	213.3	1.37×
2048	7168	1	500.0	285.5	1.75×
8192	3584	0	1394.5	752.7	1.85×
8192	3584	1	1090.9	780.1	1.40×
8192	7168	0	2534.7	1658.1	1.53×
8192	7168	1	1906.6	1017.8	1.87×

Each 1×128 tile carries 132 B in FP8 (128 B data + 4 B scale) instead of 256 B in BF16, i.e. 48.4% less activation traffic.