diff --git a/.gitignore b/.gitignore deleted file mode 100644 index 8a627a7e76..0000000000 --- a/.gitignore +++ /dev/null @@ -1,45 +0,0 @@ -.venv -*.o -*.swp -*.ii -*.ptx -*.cubin -*.fatbin* -*.module_id -*.nsys-rep -*.ncu-rep -*.sqlite -*.eggs -build/ -*.so -*.egg-info -__pycache__ -.ycm_extra_conf.py -.vimrc -.vs -.vscode -.cache -.hypothesis -.devcontainer.json -tests/cpp/build/ -.ipynb_checkpoints -*.log -CMakeFiles/CMakeSystem.cmake -sdist/ -var/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST -develop-eggs/ -dist/ -downloads/ -.pytest_cache/ -compile_commands.json -.nfs -tensor_dumps/ -artifacts/ -.DS_Store -.claude/ diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions index 2c6135a721..6add276741 160000 --- a/3rdparty/nccl-extensions +++ b/3rdparty/nccl-extensions @@ -1 +1 @@ -Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 +Subproject commit 6add276741f15619b1a3c2a7a959ee72bb7bcb59 diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 47732196ed..be6dbed3ab 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -97,6 +97,26 @@ static std::vector expected_recv_values_sorted( return vals; } +// Sorted multiset of per-token identity bytes each recv_rank should receive, +// derived from the deterministic routing map (same id_of stamp as the test). +static std::vector expected_recv_ids_sorted( + int recv_rank, int num_processes, int num_tokens, int top_k, + int num_experts, int num_local_experts) { + int base = recv_rank * num_local_experts; + std::vector ids; + for (int src = 0; src < num_processes; ++src) { + auto idx = routing_balanced(src, num_tokens, top_k, num_experts, num_local_experts); + for (int t = 0; t < num_tokens; ++t) + for (int k = 0; k < top_k; ++k) { + int64_t e = idx[t * top_k + k]; + if (e >= base && e < base + num_local_experts) + ids.push_back(static_cast((src * num_tokens + t + 1) & 0xFF)); + } + } + std::sort(ids.begin(), ids.end()); + return ids; +} + // 2^-5 relative tolerance for BF16 (matches mantissa precision with margin), // plus a small atol floor for near-zero expected values. static constexpr float kBf16Rtol = 1.0f / 32.0f; @@ -368,6 +388,123 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) { NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } +// ============================================================================= +// EPDispatchScaledTest: MXFP8 dispatch routes each token's e8m0 scale row in +// lockstep with its e4m3 data row. +// ============================================================================= + +// Each source token stamps a nonzero identity byte across its whole data row +// and its whole scale row. Dispatch is a pure permutation, so every filled recv +// slot must carry the same id in both buffers; a zero id means scales were not +// routed. +TEST_F(EpOpTestBase, MXFP8DispatchScales) { + if (g_sm_major < 9) GTEST_SKIP() << "EP requires SM_90+"; + // MXFP8 dispatch is supported on all HT EM modes (local_permute, local_dup, + // nvlink_dup); the mode is selected by the NCCL_EP_HT_EM_* env at group init. + // MXFP8 e8m0 scale bytes/token = hidden/32; NCCL EP HT requires 16B alignment. + if (hidden_dim_ % 512 != 0) + GTEST_SKIP() << "MXFP8 dispatch needs hidden_dim % 512 == 0 (16B scale alignment)"; + const int scale_cols = hidden_dim_ / 32; + using Shape = std::vector; + + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + auto h_idx = routing_balanced(g_process_id, num_tokens_, top_k_, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens_ * top_k_, 1.0f / top_k_); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + + auto id_of = [&](int tok) { + return static_cast((g_process_id * num_tokens_ + tok + 1) & 0xFF); + }; + std::vector h_data(num_tokens_ * hidden_dim_); + std::vector h_scale(num_tokens_ * scale_cols); + for (int tok = 0; tok < num_tokens_; ++tok) { + const uint8_t id = id_of(tok); + std::fill_n(&h_data[tok * hidden_dim_], hidden_dim_, id); + std::fill_n(&h_scale[tok * scale_cols], scale_cols, id); + } + + DevBuf d_data(num_tokens_ * hidden_dim_), d_scale(num_tokens_ * scale_cols); + DevBuf d_recv_data(buf.recv_capacity * hidden_dim_); + DevBuf d_recv_scale(buf.recv_capacity * scale_cols); + NVTE_CHECK_CUDA(cudaMemcpy(d_data.get(), h_data.data(), h_data.size(), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(d_scale.get(), h_scale.data(), h_scale.size(), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemset(d_recv_data.get(), 0, d_recv_data.bytes())); + NVTE_CHECK_CUDA(cudaMemset(d_recv_scale.get(), 0, d_recv_scale.bytes())); + + TensorWrapper tokens(NVTE_MXFP8_1D_SCALING); + tokens.set_rowwise_data(d_data.get(), DType::kFloat8E4M3, + Shape{(size_t)num_tokens_, (size_t)hidden_dim_}); + tokens.set_rowwise_scale_inv(d_scale.get(), DType::kFloat8E8M0, + Shape{(size_t)num_tokens_, (size_t)scale_cols}); + TensorWrapper recv_tokens(NVTE_MXFP8_1D_SCALING); + recv_tokens.set_rowwise_data(d_recv_data.get(), DType::kFloat8E4M3, + Shape{buf.recv_capacity, (size_t)hidden_dim_}); + recv_tokens.set_rowwise_scale_inv(d_recv_scale.get(), DType::kFloat8E8M0, + Shape{buf.recv_capacity, (size_t)scale_cols}); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), + t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), + tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector counts(num_local_experts_); + NVTE_CHECK_CUDA(cudaMemcpy(counts.data(), buf.recv_tokens_per_expert.get(), + num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); + auto exp_counts = expected_recv_tokens_per_expert(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + int total_recv = 0; + for (int e = 0; e < num_local_experts_; ++e) { + EXPECT_EQ(counts[e], exp_counts[e]) << "local expert " << e; + total_recv += exp_counts[e]; + } + ASSERT_LE(total_recv, static_cast(buf.recv_capacity)); + + std::vector r_data(buf.recv_capacity * hidden_dim_); + std::vector r_scale(buf.recv_capacity * scale_cols); + NVTE_CHECK_CUDA(cudaMemcpy(r_data.data(), d_recv_data.get(), r_data.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(r_scale.data(), d_recv_scale.get(), r_scale.size(), cudaMemcpyDeviceToHost)); + + std::vector got_data_ids, got_scale_ids; + got_data_ids.reserve(total_recv); + got_scale_ids.reserve(total_recv); + size_t slot = 0; + for (int e = 0; e < num_local_experts_; ++e) { + for (int i = 0; i < counts[e]; ++i, ++slot) { + got_data_ids.push_back(r_data[slot * hidden_dim_]); + got_scale_ids.push_back(r_scale[slot * scale_cols]); + } + } + std::sort(got_data_ids.begin(), got_data_ids.end()); + std::sort(got_scale_ids.begin(), got_scale_ids.end()); + + auto exp_ids = expected_recv_ids_sorted(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + ASSERT_EQ(got_data_ids.size(), exp_ids.size()); + ASSERT_EQ(got_scale_ids.size(), exp_ids.size()); + for (size_t i = 0; i < exp_ids.size(); ++i) { + EXPECT_EQ(got_data_ids[i], exp_ids[i]) << "recv data id mismatch at sorted index " << i; + EXPECT_EQ(got_scale_ids[i], exp_ids[i]) << "recv scale id mismatch at sorted index " << i; + } + + if (g_process_id == 0) + printf(" MXFP8DispatchScales: passed (recv=%d, data + scales match routing map)\n", total_recv); + + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + // ============================================================================= // EPCombineTest: round-trip identity expert -> result == top_k * tokens. // ============================================================================= diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index 7cf6017090..37f324eabb 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -7,7 +7,7 @@ /* * Shared TE EP test infrastructure. Include once per TU; ep_bootstrap() in * each test binary's main() populates process-level globals. - * Defaults: 4 experts/rank, hidden_dim=256, max_tokens_per_rank=64. + * Defaults: 4 experts/rank, hidden_dim=512, max_tokens_per_rank=64. */ #pragma once @@ -48,7 +48,7 @@ static int g_num_processes = -1; static int g_sm_major = -1; // set by ep_bootstrap; -1 until then static int g_ep_size = -1; static int g_num_experts = -1; -static int g_hidden_dim = 256; +static int g_hidden_dim = 512; static int g_max_tokens_per_rank = 64; static NVTEDType g_max_token_dtype = kNVTEFloat32; // staging-buffer sizing static bool g_ep_initialized = false; diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 4778498b7d..3c8543deee 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist +from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -19,12 +20,12 @@ ep_dispatch, ep_combine, symm_mem_alloc, + release_symm_mem_pool, is_symm_backed, _ep_combine_raw, _ep_dispatch_raw, ) - ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" OVERFLOW = os.environ.get("NVTE_EP_OVERFLOW", "0") == "1" @@ -32,11 +33,13 @@ # Must come after the transformer_engine import so libtransformer_engine.so is loaded. import transformer_engine_torch as tex # noqa: F401 - NUM_LOCAL_EXPERTS = 2 -HIDDEN_DIM = 32 +# MXFP8 dispatch needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0. Defaults +# satisfy both so the MXFP8 tests run by default; override via NVTE_EP_HIDDEN_DIM / +# NVTE_EP_TOKENS_PER_RANK. +HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512")) TOP_K = 2 -TOKENS_PER_RANK = 4 +TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32")) def _zero_copy_test_include(fn): @@ -57,6 +60,18 @@ def _overflow_test_include(fn): return fn +# MXFP8 grouped dispatch needs a per-expert alignment of 128, but the EP backend caches a single +# alignment per process, so alignment=128 tests cannot share a process with the alignment=0 tests. +# They run in a dedicated pass (NVTE_EP_MXFP8_PASS=1) instead. +MXFP8_PASS = os.environ.get("NVTE_EP_MXFP8_PASS", "0") == "1" + + +def _mxfp8_align_test(fn): + """Mark a test that dispatches with alignment=128; runs only in the MXFP8 pass.""" + fn._mxfp8_align_test = True + return fn + + class _StageToSymm(torch.autograd.Function): """Identity op that stages ``src`` into a symm-mem buffer; grad passes through. Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine. @@ -121,6 +136,16 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _degroup_mxfp8(recv_grouped, valid_counts=None): + """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. + With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert + slot; otherwise return every (padded) row.""" + parts = recv_grouped.split_into_quantized_tensors() + if valid_counts is None: + return torch.cat([p.dequantize() for p in parts], dim=0) + return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) + + class _Cfg: rank: int world_size: int @@ -171,6 +196,13 @@ def setUpClass(cls): ) def setUp(self): + # alignment=128 MXFP8 tests run only in the dedicated MXFP8 pass; everything else skips + # there (and the MXFP8 tests skip outside it) since the backend pins one alignment/process. + is_mxfp8_align = getattr(getattr(self, self._testMethodName), "_mxfp8_align_test", False) + if MXFP8_PASS and not is_mxfp8_align: + self.skipTest("only alignment=128 MXFP8 tests run in the MXFP8 pass") + if not MXFP8_PASS and is_mxfp8_align: + self.skipTest("alignment=128 MXFP8 tests run in the dedicated MXFP8 pass") # Only the zero-copy-capable tests run in the zero-copy pass. if ZERO_COPY and not getattr( getattr(self, self._testMethodName), "_zero_copy_test_include", False @@ -185,7 +217,13 @@ def setUp(self): ): self.skipTest("not exercised in overflow mode") - def _make_buffer(self, alignment=0, top_k=TOP_K): + def _make_buffer( + self, + alignment=0, + top_k=TOP_K, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, + ): return EpBuffer( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, @@ -193,6 +231,8 @@ def _make_buffer(self, alignment=0, top_k=TOP_K): num_local_experts=NUM_LOCAL_EXPERTS, recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, alignment=alignment, + dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, + combine_bwd_quant_recipe=combine_bwd_quant_recipe, ) def _expert_out(self, expert_out): @@ -335,6 +375,70 @@ def test_dispatch_autograd(self): tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2 ) + # MXFP8 dispatch + + def _mxfp8_quantizer(self): + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + + def _require_mxfp8_shapes(self): + if HIDDEN_DIM % 512 != 0 or TOKENS_PER_RANK % 32 != 0: + self.skipTest( + "MXFP8 needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0 " + "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)" + ) + + def _assert_mxfp8_matches_bf16(self, recv_mx, tokens, topk_idx, w, tc): + """Dequantized MXFP8 recv matches a bf16 dispatch of the same tokens. Both share the + alignment=128 padded expert-major layout, so compare the full prefix [0:sum(padded)].""" + ref_tokens = self._mxfp8_quantizer().quantize(tokens).dequantize() + ref_recv, _rw, _tc = ep_dispatch(self._make_buffer(alignment=128), ref_tokens, topk_idx, w) + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(recv_mx).float(), ref_recv.float()[:n], atol=1e-2, rtol=1e-2 + ) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_dispatch_mxfp8(self): + """MXFP8 dispatch quantizes bf16 tokens internally; recv (a per-expert GroupedTensor) + dequantized matches a bf16 dispatch of the same tokens. Under zero-copy the recv data and + scales are symm-mem backed.""" + self._require_mxfp8_shapes() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + if ZERO_COPY: + self.assertTrue(is_symm_backed(recv_mx.rowwise_data)) + self.assertTrue(is_symm_backed(recv_mx.scale_inv)) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + + @_zero_copy_test_include + @_mxfp8_align_test + def test_caller_provides_dispatch_recv_mxfp8(self): + """One caller-supplied buffer holds the recv data followed by the e8m0 scales; ep_dispatch + slices it and the returned GroupedTensor views the data and scale regions of that buffer.""" + self._require_mxfp8_shapes() + from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE + + rc = self.cfg.recv_capacity_per_rank + cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE + nbytes = rc * (HIDDEN_DIM + cols) # fp8 data + e8m0 scales, one byte per element + if ZERO_COPY: + recv_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + else: + recv_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w, recv_tokens=recv_buf) + # the returned GroupedTensor views the caller buffer's data then scale regions + self.assertEqual(recv_mx.rowwise_data.data_ptr(), recv_buf.data_ptr()) + self.assertEqual(recv_mx.scale_inv.data_ptr(), recv_buf.data_ptr() + rc * HIDDEN_DIM) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + @_zero_copy_test_include def test_caller_provides_dispatch_recv_tokens(self): """Caller-supplied recv_tokens (symm-mem-backed in zero-copy): ep_dispatch @@ -382,6 +486,48 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) + @_mxfp8_align_test + def test_combine_bwd_mxfp8_caller_grad_out(self): + """MXFP8 combine backward into a single caller buffer sliced into data + e8m0 scales: the + returned per-expert GroupedTensor views those regions and, dequantized, matches a bf16 + combine backward reference on the same routing. MXFP8 combine backward runs in the + non-zero-copy path, so this test skips the zero-copy pass.""" + if ZERO_COPY: + self.skipTest("MXFP8 combine backward is not supported under zero-copy") + self._require_mxfp8_shapes() + from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE + + rc = self.cfg.recv_capacity_per_rank + cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + eo_vals = ( + torch.linspace(-0.5, 0.5, rc * HIDDEN_DIM, device=self.cfg.device) + .reshape(rc, HIDDEN_DIM) + .to(torch.bfloat16) + ) + # MXFP8 combine backward writes into one caller buffer (data then e8m0 scales) + buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w) # seeds the routing + nbytes = rc * (HIDDEN_DIM + cols) + grad_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + src_mx = eo_vals.detach().clone().requires_grad_(True) + out_mx = ep_combine(buf_mx, src_mx, grad_out=grad_buf) + (0.5 * (out_mx.float() ** 2).sum()).backward() + g_mx = src_mx.grad # per-expert GroupedTensor viewing grad_buf + self.assertEqual(g_mx.rowwise_data.data_ptr(), grad_buf.data_ptr()) + self.assertEqual(g_mx.scale_inv.data_ptr(), grad_buf.data_ptr() + rc * HIDDEN_DIM) + # bf16 reference combine backward on the same routing + buf_bf = self._make_buffer(alignment=128) + ep_dispatch(buf_bf, tokens, topk_idx, w) + src_bf = eo_vals.detach().clone().requires_grad_(True) + out_bf = ep_combine(buf_bf, src_bf) + (0.5 * (out_bf.float() ** 2).sum()).backward() + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + ) + @_zero_copy_test_include def test_zero_copy_pool_auto_alloc(self): """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO @@ -594,5 +740,7 @@ def _init_distributed(): result = runner.run(suite) dist.barrier() ep_finalize() + # Deregister symm-mem windows while the comm is still valid. + release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 10e814eb80..5788c2ac9f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -47,11 +47,13 @@ run_pass() { local zc="$2" local eager="${3:-0}" local overflow="${4:-0}" + local mxfp8="${5:-0}" local log="stdout_ep_${label}.txt" echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun. - NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" setsid timeout --foreground \ - --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ + NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" \ + NVTE_EP_MXFP8_PASS="${mxfp8}" \ + setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ "${SCRIPT}" 2>&1 | tee "${log}" local rc=${PIPESTATUS[0]} @@ -72,5 +74,11 @@ run_pass "default" 0 run_pass "zero_copy" 1 run_pass "eager" 0 1 run_pass "overflow" 0 0 1 +# MXFP8 grouped dispatch pins the per-expert alignment to 128, which the backend caches +# process-wide, so its tests get their own passes (normal + zero-copy + eager IO). mxfp8 is the +# 5th arg; eager is the 3rd. +run_pass "mxfp8" 0 0 0 1 +run_pass "mxfp8_zero_copy" 1 0 0 1 +run_pass "mxfp8_eager" 0 1 0 1 exit $RET diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index dd44773514..a5f4ad4454 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -47,26 +47,50 @@ ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { return ncclFloat8e4m3; case kNVTEFloat8E5M2: return ncclFloat8e5m2; + case kNVTEFloat8E8M0: + return ncclUint8; default: NVTE_ERROR("Unsupported NVTEDType for NCCL dtype conversion: ", static_cast(dtype)); } return ncclFloat32; // unreachable } -// shape_out is caller-owned; desc.sizes aliases shape_out.data and must -// outlive the NCCL EP call. +// Which part of a TE tensor an NCCL descriptor points at: the data payload, or +// (for block-scaled tensors) the rowwise scale-inverse that rides alongside it. +enum class DescSource { kData, kScaleInv }; + +// Build an NCCL descriptor for a TE tensor's data (kData) or its rowwise +// scale-inverse (kScaleInv). shape_out is caller-owned; desc.sizes aliases +// shape_out.data and must outlive the NCCL EP call. Uses the matching window +// field (win.window / win.scale_window) when set, else the raw pointer. inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, NVTEShape& shape_out, - const NVTECommWindow& win = {}) { - shape_out = nvte_tensor_shape(t); + const NVTECommWindow& win = {}, + DescSource source = DescSource::kData) { ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; + void* raw_ptr = nullptr; + ncclWindow_t win_hdl = nullptr; + uint64_t win_offset = 0; + if (source == DescSource::kData) { + shape_out = nvte_tensor_shape(t); + desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); + raw_ptr = nvte_tensor_data(t); + win_hdl = win.window; + win_offset = win.offset; + } else { + const SimpleTensor& si = convertNVTETensorCheck(t)->scale_inv; + shape_out = nvte_make_shape(si.shape.data(), si.shape.size()); + desc.datatype = te_dtype_to_nccl_dtype(static_cast(si.dtype)); + raw_ptr = si.dptr; + win_hdl = win.scale_window; + win_offset = win.scale_offset; + } desc.ndim = shape_out.ndim; desc.sizes = shape_out.data; - desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); - if (win.window != nullptr) { - desc.win_hdl = win.window; - desc.win_offset = win.offset; + if (win_hdl != nullptr) { + desc.win_hdl = win_hdl; + desc.win_offset = win_offset; } else { - desc.data = nvte_tensor_data(t); + desc.data = raw_ptr; NVTE_CHECK(desc.data != nullptr, "tensor data must not be null"); } return desc; @@ -411,16 +435,45 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, recv_topk_weights_win); } + // Block-scaled (e.g. MXFP8): route the per-token scale-inverse alongside the + // data. High-precision (bf16/fp16/fp32) and per-tensor FP8 payloads carry the + // default delayed scaling mode and skip this. Keys on is_block_scaling so + // NVFP4 can reuse this path later. + const NVTEScalingMode tokens_scaling_mode = nvte_tensor_scaling_mode(tokens); + const bool is_scaled = is_block_scaling(tokens_scaling_mode); + NVTEShape scales_in_shape, scales_out_shape; + ncclEpTensor_t nccl_scales_in = NCCL_EP_TENSOR_INIT, nccl_scales_out = NCCL_EP_TENSOR_INIT; + if (is_scaled) { + NVTE_CHECK(is_mxfp8_scaling(tokens_scaling_mode), + "EP dispatch supports MXFP8 block scaling only; got scaling mode ", + static_cast(tokens_scaling_mode)); + NVTE_CHECK(nvte_tensor_scaling_mode(recv_tokens) == tokens_scaling_mode, + "recv_tokens scaling mode must match tokens scaling mode"); + nccl_scales_in = + make_nccl_ep_tensor(tokens, scales_in_shape, tokens_win, DescSource::kScaleInv); + nccl_scales_out = + make_nccl_ep_tensor(recv_tokens, scales_out_shape, recv_tokens_win, DescSource::kScaleInv); + } else { + NVTE_CHECK(!is_fp8_dtype(static_cast(tok_dtype)), + "EP dispatch of FP8 tokens requires a block scaling mode (e.g. MXFP8); " + "per-tensor (delayed) FP8 scaling is not supported"); + } + ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; in_struct.tokens = &nccl_tokens_in; in_struct.topk_weights = is_forward ? &nccl_topk_weights_in : nullptr; + in_struct.scales = is_scaled ? &nccl_scales_in : nullptr; ncclEpDispatchOutputs_t out_struct = NCCL_EP_DISPATCH_OUTPUTS_INIT; out_struct.tokens = &nccl_tokens_out; out_struct.topk_weights = is_forward ? &nccl_topk_weights_out : nullptr; + out_struct.scales = is_scaled ? &nccl_scales_out : nullptr; ncclEpDispatchConfig_t dispatch_cfg = NCCL_EP_DISPATCH_CONFIG_INIT; dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; + // Block-scaled payloads forward the per-token scale-inverse; select the matching recipe. + dispatch_cfg.quantization_recipe = + is_scaled ? NCCL_EP_DISPATCH_QUANT_SCALES_FORWARD : NCCL_EP_DISPATCH_QUANT_NONE; std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h index 424c350bbd..ed67774493 100644 --- a/transformer_engine/common/include/transformer_engine/comm_window.h +++ b/transformer_engine/common/include/transformer_engine/comm_window.h @@ -26,6 +26,9 @@ struct ncclWindow_vidmem; typedef struct { struct ncclWindow_vidmem* window; /*!< NCCL window, or NULL to use the raw data pointer. */ uint64_t offset; /*!< Byte offset of the payload within window. */ + struct ncclWindow_vidmem* + scale_window; /*!< Window for a block-scaled tensor's scale-inverse, or NULL for raw. */ + uint64_t scale_offset; /*!< Byte offset of the scale-inverse within scale_window. */ } NVTECommWindow; #ifdef __cplusplus diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index a4800944b5..31417fbff1 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -150,6 +150,12 @@ void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv * *_win arguments enable zero-copy via symmem windows; pass NVTECommWindow{} * when unused. Requires a prior nvte_ep_prepare on this handle_mem. * + * tokens/recv_tokens may be high-precision (bf16/fp16) or FP8: + * for the latter, set rowwise data and rowwise scale-inverse (unswizzled + * [T, hidden_dim/block]) on the tensor and the scales are routed alongside the + * data. tokens and recv_tokens must share a scaling mode. For now, only MXFP8 + * (NVTE_MXFP8_1D_SCALING, e4m3 data + e8m0 scales) is supported. + * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] topk_idx [T, top_k] int64 sparse routing indices. * \param[in] tokens [T, hidden_dim] input tokens. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 95f02c64f0..c952a738eb 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -691,14 +691,18 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_pe at::Tensor total_recv_tokens); void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, - at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights); + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights, + std::optional tokens_scale_inv = std::nullopt, + std::optional recv_scale_inv = std::nullopt); void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result); void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights, at::Tensor grad_tokens, at::Tensor grad_topk_weights); -void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out); +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out, + std::optional grad_scale_inv = std::nullopt, + std::optional grad_expert_out_scale_inv = std::nullopt); // Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP. void register_ep_bindings(pybind11::module_ &m); diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 9bf821721f..cef489dbbb 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -73,8 +73,14 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) { NVTE_CHECK(nccl_sm != nullptr, "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to " "\"NCCL\" before allocating EP payload buffers."); - return NVTECommWindow{static_cast(nccl_sm->get_window()), - static_cast(nccl_sm->get_offset())}; + // rendezvous resolves ``t`` by its storage base, so get_offset() is the allocation's offset in + // the NCCL window. Add ``t``'s own storage offset so a slice/view of a symm-mem allocation + // (e.g. the scale region carved from a shared recv buffer) resolves to its true position in the + // window rather than the allocation base. + const uint64_t offset = + static_cast(nccl_sm->get_offset()) + + static_cast(t.storage_offset()) * static_cast(t.element_size()); + return NVTECommWindow{static_cast(nccl_sm->get_window()), offset}; #else (void)t; return kNoWindow; @@ -114,6 +120,34 @@ DType check_topk_idx_dtype(at::Tensor topk_idx) { using Shape = std::vector; +// EP block scaling supports only E4M3 MXFP8 today. A future block-scaled recipe (e.g. NVFP4) +// would also set is_scaled but carry a non-FP8 token dtype, so key the guard on the FP8 dtype. +bool is_mxfp8_scaled(bool is_scaled, const at::Tensor& data) { + return is_scaled && data.scalar_type() == at::kFloat8_e4m3fn; +} + +// Validate the scale-inverse pair of a block-scaled EP op and return the scale column count +// (hidden/block). Both scales must be 2D contiguous with matching cols dividing H and numels +// equal to their row counts times cols; the recv scale must be symm-mem-backed under zero-copy. +size_t check_mxfp8_scale_pair(const at::Tensor& send_scale, const at::Tensor& recv_scale, + size_t send_rows, size_t recv_rows, size_t H, const char* recv_name) { + NVTE_CHECK(send_scale.dim() >= 2 && recv_scale.dim() >= 2, + "scale-inverses must be at least 2D [., H/block]"); + NVTE_CHECK(send_scale.is_contiguous() && recv_scale.is_contiguous(), + "scale-inverses must be contiguous"); + const size_t sc_cols = static_cast(send_scale.size(-1)); + NVTE_CHECK(sc_cols > 0 && H % sc_cols == 0, "scale cols (", sc_cols, + ") must be a non-zero divisor of hidden (", H, ")"); + NVTE_CHECK(static_cast(recv_scale.size(-1)) == sc_cols, + "recv scale cols must match send scale cols"); + NVTE_CHECK(static_cast(send_scale.numel()) == send_rows * sc_cols, + "send scale numel must equal rows * cols"); + NVTE_CHECK(static_cast(recv_scale.numel()) == recv_rows * sc_cols, + "recv scale numel must equal rows * cols"); + check_symm_mem_required(recv_scale, recv_name); + return sc_cols; +} + } // namespace bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_relaxed); } @@ -218,8 +252,13 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_pe total_recv_tokens_te.data(), &layer_cfg, stream); } +// tokens_scale_inv / recv_scale_inv are set only for block-scaled dispatch (for +// now MXFP8): tokens/recv_tokens carry e4m3 data and the scale tensors carry the +// unswizzled e8m0 scale-inverses [T, H/block]. Both null => bf16/fp16/fp32. void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, - at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights) { + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights, + std::optional tokens_scale_inv, + std::optional recv_scale_inv) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(tokens.dim() >= 2, "tokens must be at least 2D [..., H]"); NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); @@ -250,16 +289,39 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, check_symm_mem_required(recv_tokens, "recv_tokens"); check_symm_mem_required(recv_topk_weights, "recv_topk_weights"); + // Block-scaled dispatch: tokens carry e4m3 data and the scale tensors carry + // unswizzled e8m0 scale-inverses [T, H/block]. Scales ride in the tensor; the + // backend keys on the tensor's scaling mode. is_mxfp8 is split from is_scaled + // so future block-scaled recipes can reuse the scale-routing plumbing while + // building their own TE tensors; only MXFP8 is supported for now. + const bool is_scaled = tokens_scale_inv.has_value(); + const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, tokens); + size_t sc_cols = 0; + if (is_scaled) { + NVTE_CHECK(recv_scale_inv.has_value(), + "recv_scale_inv must be provided together with tokens_scale_inv"); + NVTE_CHECK(is_mxfp8, "EP dispatch currently supports only E4M3 MXFP8 block scaling"); + sc_cols = check_mxfp8_scale_pair(*tokens_scale_inv, *recv_scale_inv, T_flat, recv_pr, H, + "recv_scale_inv"); + } + auto tok_dtype = GetTransformerEngineDType(tokens.scalar_type()); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); auto topk_idx_te = makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, idx_dtype); - auto tokens_te = makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); + auto tokens_te = + is_mxfp8 ? makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype, + nullptr, nullptr, tokens_scale_inv->data_ptr(), + Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); auto topk_w_te = makeTransformerEngineTensor(topk_weights.data_ptr(), Shape{T_flat, topk_n}, DType::kFloat32); auto recv_tokens_te = - makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); + is_mxfp8 ? makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype, + nullptr, nullptr, recv_scale_inv->data_ptr(), + Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); auto recv_topk_w_te = makeTransformerEngineTensor(recv_topk_weights.data_ptr(), Shape{recv_pr}, DType::kFloat32); @@ -269,6 +331,17 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, NVTECommWindow topk_w_win = maybe_make_window(topk_weights); NVTECommWindow recv_tokens_win = maybe_make_window(recv_tokens); NVTECommWindow recv_topk_w_win = maybe_make_window(recv_topk_weights); + // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window. + // Send scales (tokens_scale_inv) stay staged like the send data; recv scales + // are the one-sided write target and must be symm-mem-backed under zero-copy. + if (is_scaled) { + const NVTECommWindow tsi_win = maybe_make_window(*tokens_scale_inv); + const NVTECommWindow rsi_win = maybe_make_window(*recv_scale_inv); + tokens_win.scale_window = tsi_win.window; + tokens_win.scale_offset = tsi_win.offset; + recv_tokens_win.scale_window = rsi_win.window; + recv_tokens_win.scale_offset = rsi_win.offset; + } nvte_ep_dispatch(handle_mem_te.data(), topk_idx_te.data(), tokens_te.data(), tokens_win, topk_w_te.data(), topk_w_win, recv_tokens_te.data(), recv_tokens_win, recv_topk_w_te.data(), recv_topk_w_win, stream); @@ -344,7 +417,9 @@ void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_t g_recv_w_win, grad_tokens_te.data(), grad_topk_w_te.data(), stream); } -void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out) { +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out, + std::optional grad_scale_inv, + std::optional grad_expert_out_scale_inv) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(grad.dim() >= 2, "grad must be at least 2D [..., H]"); NVTE_CHECK(grad_expert_out.dim() >= 2, "grad_expert_out must be at least 2D [..., recv_pr, H]"); @@ -363,17 +438,51 @@ void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expe // EpBuffer-owned scatter target and must be symm-mem in zero-copy mode. check_symm_mem_required(grad_expert_out, "grad_expert_out"); + // Block-scaled (MXFP8) backward: grad/grad_expert_out carry e4m3 data and the scale + // tensors carry the unswizzled e8m0 scale-inverses [., H/block]; the reverse-direction + // dispatch forwards them like the forward path. is_mxfp8 is split from is_scaled so + // future block-scaled recipes can reuse the plumbing; only MXFP8 is supported for now. + const bool is_scaled = grad_scale_inv.has_value(); + const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, grad); + size_t sc_cols = 0; + if (is_scaled) { + NVTE_CHECK(grad_expert_out_scale_inv.has_value(), + "grad_expert_out_scale_inv must be provided together with grad_scale_inv"); + NVTE_CHECK(is_mxfp8, "EP combine backward currently supports only E4M3 MXFP8 block scaling"); + sc_cols = check_mxfp8_scale_pair(*grad_scale_inv, *grad_expert_out_scale_inv, T_flat, recv_pr, + H, "grad_expert_out_scale_inv"); + } + auto g_dtype = GetTransformerEngineDType(grad.scalar_type()); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); - auto grad_te = makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); + auto grad_te = is_mxfp8 + ? makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype, + nullptr, nullptr, grad_scale_inv->data_ptr(), + Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); auto grad_expert_out_te = - makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); + is_mxfp8 + ? makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype, + nullptr, nullptr, grad_expert_out_scale_inv->data_ptr(), + Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); // grad is autograd-allocated (staged); grad_expert_out resolves to a symm-mem // window in zero-copy mode, else kNoWindow for the staged path. NVTECommWindow grad_win = maybe_make_window(grad); NVTECommWindow grad_expert_out_win = maybe_make_window(grad_expert_out); + // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window, + // mirroring the forward dispatch. Send scales stay staged like the send data; + // recv scales are the one-sided write target and must be symm-mem-backed. + if (is_scaled) { + const NVTECommWindow gsi_win = maybe_make_window(*grad_scale_inv); + const NVTECommWindow gesi_win = maybe_make_window(*grad_expert_out_scale_inv); + grad_win.scale_window = gsi_win.window; + grad_win.scale_offset = gsi_win.offset; + grad_expert_out_win.scale_window = gesi_win.window; + grad_expert_out_win.scale_offset = gesi_win.offset; + } nvte_ep_combine_bwd(handle_mem_te.data(), grad_te.data(), grad_win, grad_expert_out_te.data(), grad_expert_out_win, stream); } @@ -396,11 +505,16 @@ void register_ep_bindings(pybind11::module_& m) { py::arg("tokens_per_expert"), py::arg("top_k"), py::arg("dispatch_output_per_expert_alignment"), py::arg("total_recv_tokens"), py::call_guard()); - m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::call_guard()); + m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::arg("handle_mem"), py::arg("topk_idx"), + py::arg("tokens"), py::arg("topk_weights"), py::arg("recv_tokens"), + py::arg("recv_topk_weights"), py::arg("tokens_scale_inv") = std::nullopt, + py::arg("recv_scale_inv") = std::nullopt, py::call_guard()); m.def("ep_combine", &ep_combine, "EP combine", py::call_guard()); m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward", py::call_guard()); - m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", + m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", py::arg("handle_mem"), + py::arg("grad"), py::arg("grad_expert_out"), py::arg("grad_scale_inv") = std::nullopt, + py::arg("grad_expert_out_scale_inv") = std::nullopt, py::call_guard()); } diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d1525b53f0..647e4830a5 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Methods needed for distributed training (DP/TP).""" + from __future__ import annotations from collections.abc import Iterable @@ -50,7 +51,6 @@ from .tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..debug.pytorch.debug_quantization import DebugQuantizedTensor - __all__ = ["checkpoint", "CudaRNGStatesTracker"] @@ -1882,6 +1882,11 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr _SYMM_MEM_POOL = None _SYMM_MEM_POOL_BACKEND = None +# Device the pool was created for; the torch symm_mem._symm_mem_pools cache is keyed by it. +_SYMM_MEM_POOL_DEVICE = None +# True when the pool was created via torch's get_mem_pool, which caches it in the private +# symm_mem._symm_mem_pools dict; release then has to drop that cached reference. +_SYMM_MEM_POOL_TORCH_CACHED = False def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): @@ -1890,12 +1895,14 @@ def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): backend arg, so the (process-global) backend is always set before the pool is created. The collective rendezvous cost is amortized across allocations (paid per new segment, not per buffer). """ - global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED if _SYMM_MEM_POOL is None: symm_mem.set_backend(backend) _SYMM_MEM_POOL_BACKEND = backend + _SYMM_MEM_POOL_DEVICE = device if hasattr(symm_mem, "get_mem_pool"): _SYMM_MEM_POOL = symm_mem.get_mem_pool(device) + _SYMM_MEM_POOL_TORCH_CACHED = True elif hasattr(torch.cuda, "MemPool") and hasattr(symm_mem, "get_mempool_allocator"): _SYMM_MEM_POOL = torch.cuda.MemPool(symm_mem.get_mempool_allocator(device)) else: @@ -1911,6 +1918,41 @@ def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): return _SYMM_MEM_POOL +def release_symm_mem_pool(device: Optional[torch.device] = None) -> None: + """Free the process-wide symm-mem pool's segments, deregistering their NCCL windows. + + Call before ``dist.destroy_process_group()``: the pool's windows are registered on + the group's NCCL comm, which becomes invalid once the group is destroyed. No-op if + no pool was created. ``device`` is ignored; the pool's creation device is used. + """ + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED + if _SYMM_MEM_POOL is None: + return + # Key the cache lookup on the device the pool was created for, not the caller's current + # device, so teardown works even when they differ. + device = _SYMM_MEM_POOL_DEVICE + _SYMM_MEM_POOL = None + _SYMM_MEM_POOL_BACKEND = None + _SYMM_MEM_POOL_DEVICE = None + torch_cached = _SYMM_MEM_POOL_TORCH_CACHED + _SYMM_MEM_POOL_TORCH_CACHED = False + # A pool from torch's get_mem_pool is also cached in the private module dict + # symm_mem._symm_mem_pools; drop that reference so the segments' refcount reaches zero + # and their NCCL windows deregister. Fail loudly if this internal has changed shape, + # otherwise the windows would silently leak past destroy_process_group(). + if torch_cached: + pools = getattr(symm_mem, "_symm_mem_pools", None) + if not isinstance(pools, dict) or device not in pools: + raise RuntimeError( + "torch symmetric-memory pool cache (symm_mem._symm_mem_pools) is missing or " + "has changed layout; cannot release the pooled segments and their NCCL windows " + "would leak past destroy_process_group(). This torch version needs an updated " + "release_symm_mem_pool()." + ) + pools.pop(device, None) + torch.cuda.empty_cache() + + def symm_mem_alloc( shape, dtype: torch.dtype, diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 940812b6a4..d70f7b3fcc 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -7,7 +7,7 @@ import atexit import warnings -from typing import Optional +from typing import Optional, TYPE_CHECKING import torch import torch.distributed as dist @@ -15,8 +15,12 @@ import transformer_engine_torch as tex from .cpu_offload import mark_not_offload -from .distributed import symm_mem_alloc +from .distributed import symm_mem_alloc, release_symm_mem_pool +# Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of +# common.recipe (the concrete recipe classes are imported lazily where used). +if TYPE_CHECKING: + from ..common.recipe import Recipe __all__ = [ "EpBuffer", @@ -25,6 +29,7 @@ "ep_dispatch", "ep_combine", "symm_mem_alloc", + "release_symm_mem_pool", "is_symm_backed", ] @@ -231,6 +236,9 @@ class EpBuffer: "eager", "total_recv_tokens", "_host_total_recv_tokens", + "dispatch_fwd_quant_recipe", + "combine_bwd_quant_recipe", + "prepared", ) def __init__( @@ -243,6 +251,8 @@ def __init__( alignment: int = 0, payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, + dispatch_fwd_quant_recipe: Optional["Recipe"] = None, + combine_bwd_quant_recipe: Optional["Recipe"] = None, ) -> None: if not _BOOTSTRAPPED: raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") @@ -268,6 +278,8 @@ def __init__( self.payload_dtype = payload_dtype self.device = device self.zero_copy = bool(tex.ep_get_zero_copy()) + self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe + self.combine_bwd_quant_recipe = combine_bwd_quant_recipe size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) @@ -283,6 +295,8 @@ def __init__( mark_not_offload(self.total_recv_tokens) # Host mirror of total_recv_tokens, set by ep_prepare in eager mode. self._host_total_recv_tokens: Optional[int] = None + # Set by ep_prepare; dispatch/combine read the routing it cached in handle_mem. + self.prepared = False # torch.library custom ops (so they don't graph-break under torch.compile) @@ -313,7 +327,7 @@ def _(*_args, **_kw): @torch.library.custom_op( f"{_LIB}::dispatch", - mutates_args=("recv_tokens", "recv_topk_weights"), + mutates_args=("recv_tokens", "recv_topk_weights", "recv_scale_inv"), device_types="cuda", ) def _dispatch_op( @@ -323,8 +337,19 @@ def _dispatch_op( topk_weights: torch.Tensor, recv_tokens: torch.Tensor, recv_topk_weights: torch.Tensor, + tokens_scale_inv: Optional[torch.Tensor] = None, + recv_scale_inv: Optional[torch.Tensor] = None, ) -> None: - tex.ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights) + tex.ep_dispatch( + handle_mem, + topk_idx, + tokens, + topk_weights, + recv_tokens, + recv_topk_weights, + tokens_scale_inv, + recv_scale_inv, + ) @_dispatch_op.register_fake @@ -372,15 +397,17 @@ def _(*_args, **_kw): @torch.library.custom_op( f"{_LIB}::combine_bwd", - mutates_args=("grad_expert_out",), + mutates_args=("grad_expert_out", "grad_expert_out_scale_inv"), device_types="cuda", ) def _combine_bwd_op( handle_mem: torch.Tensor, grad: torch.Tensor, grad_expert_out: torch.Tensor, + grad_scale_inv: Optional[torch.Tensor] = None, + grad_expert_out_scale_inv: Optional[torch.Tensor] = None, ) -> None: - tex.ep_combine_bwd(handle_mem, grad, grad_expert_out) + tex.ep_combine_bwd(handle_mem, grad, grad_expert_out, grad_scale_inv, grad_expert_out_scale_inv) @_combine_bwd_op.register_fake @@ -410,6 +437,7 @@ def ep_prepare(buffer: "EpBuffer", topk_idx: torch.Tensor) -> torch.Tensor: ) if buffer.eager: buffer._host_total_recv_tokens = int(buffer.total_recv_tokens.item()) + buffer.prepared = True return buffer.tokens_per_expert @@ -422,6 +450,8 @@ def _ep_dispatch_raw( recv_topk_weights: torch.Tensor, ) -> None: """Raw dispatch; no autograd, no prepare. Caller must run ep_prepare first.""" + if not buffer.prepared: + raise RuntimeError("ep_dispatch requires ep_prepare to have run on this buffer first.") tex.ep_dispatch( buffer.handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights ) @@ -429,6 +459,8 @@ def _ep_dispatch_raw( def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch.Tensor) -> None: """Raw combine; no autograd. Caller pre-weights expert_out.""" + if not buffer.prepared: + raise RuntimeError("ep_combine requires ep_prepare to have run on this buffer first.") tex.ep_combine(buffer.handle_mem, expert_out, result) @@ -442,32 +474,86 @@ class _EpDispatch(torch.autograd.Function): def forward( # type: ignore[override] ctx, handle_mem: torch.Tensor, - recv_tokens: torch.Tensor, - recv_topk_weights: torch.Tensor, + recv_tokens: Optional[torch.Tensor], + recv_topk_weights: Optional[torch.Tensor], topk_idx: torch.Tensor, tokens: torch.Tensor, topk_weights: torch.Tensor, + tokens_scale_inv: Optional[torch.Tensor] = None, + token_counts: Optional[torch.Tensor] = None, + rows: int = 0, + payload_dtype: torch.dtype = torch.bfloat16, ): - """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand.""" + """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand. When scales are set + (MXFP8 for now), ``tokens`` is the quantized tensor kept as the autograd operand so grad + reaches the pre-quant input. Recv outputs are carved/allocated here: a caller may supply + ``recv_tokens`` / ``recv_topk_weights``, else they are sized to ``rows``.""" + from .quantized_tensor import QuantizedTensor + + is_scaled = tokens_scale_inv is not None + tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens + assert tokens_data.dim() == 2, "EP dispatch tokens must be 2D [num_tokens, hidden]" + hidden = tokens_data.shape[-1] + device = tokens_data.device + zero_copy = tex.ep_get_zero_copy() + + recv_scale_inv = None + if is_scaled: + if tokens._fp8_dtype != tex.DType.kFloat8E4M3: + raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.") + # recv data + scales share one buffer (data then scales); carve or allocate it here. + recv_tokens, recv_scale_inv = _scale_alloc_io( + recv_tokens, + rows, + hidden, + tokens_scale_inv.shape[-1], + tokens_data.dtype, + tokens_scale_inv.dtype, + device, + zero_copy, + ) + # Reinterpret byte-backed FP8 data as the fp8 dtype so the backend sees a scaled tensor. + dispatch_tokens = tokens_data.view(torch.float8_e4m3fn) + dispatch_recv = recv_tokens.view(torch.float8_e4m3fn) + else: + if recv_tokens is None: + recv_tokens = _alloc_io((rows, hidden), payload_dtype, device, zero_copy) + dispatch_tokens = tokens_data + dispatch_recv = recv_tokens + if recv_topk_weights is None: + recv_topk_weights = _alloc_io((rows,), torch.float32, device, zero_copy) torch.ops.transformer_engine_ep.dispatch( handle_mem, topk_idx, - tokens, + dispatch_tokens, topk_weights, - recv_tokens, + dispatch_recv, recv_topk_weights, + tokens_scale_inv, + recv_scale_inv, ) ctx.save_for_backward(handle_mem) ctx.tokens_shape = tokens.shape - ctx.tokens_dtype = tokens.dtype ctx.topk_weights_shape = topk_weights.shape - ctx.tokens_T_flat = tokens.numel() // tokens.shape[-1] + ctx.num_tokens = tokens_data.shape[0] ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] ctx.top_k = topk_weights.shape[-1] - ctx.hidden_dim = tokens.shape[-1] + ctx.hidden_dim = hidden # Detach so the long-lived buffers aren't tracked as differentiable outputs; - # autograd re-attaches grad_fn pointing back at this Function. - return recv_tokens.detach(), recv_topk_weights.detach() + # autograd re-attaches grad_fn pointing back at this Function. For scaled inputs + # the expert-major recv data + scales are wrapped into a per-expert GroupedTensor + # so downstream grouped GEMM and autograd see a proper quantized grouped tensor. + if is_scaled: + recv_out = _make_grouped_mxfp8( + recv_tokens.view(tokens._rowwise_data.dtype), + recv_scale_inv, + token_counts, + tokens._fp8_dtype, + tokens.dtype, + ) + else: + recv_out = recv_tokens.detach() + return recv_out, recv_topk_weights.detach() @staticmethod def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] @@ -476,8 +562,10 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] device = handle_mem.device g_recv_tokens = g_recv_tokens.contiguous() g_recv_topk_weights = g_recv_topk_weights.contiguous() + # Dispatch grad follows the recv grad's (high-precision) dtype; the quantizer's STE + # owns the fp8 boundary for scaled inputs. grad_tokens = torch.empty( - ctx.tokens_T_flat, ctx.hidden_dim, dtype=ctx.tokens_dtype, device=device + ctx.num_tokens, ctx.hidden_dim, dtype=g_recv_tokens.dtype, device=device ) grad_topk_weights = torch.empty( ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device @@ -496,21 +584,21 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] None, # topk_idx grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape), + None, # tokens_scale_inv (scales; non-differentiable) + None, # token_counts (per-expert counts; non-differentiable) + None, # rows (sizing scalar) + None, # payload_dtype (sizing scalar) ) class _EpCombine(torch.autograd.Function): - """Autograd combine. - - bwd scatters the expert_out grad into ``grad_out``. When the caller supplies it - (mcore-managed mode) that buffer is used as-is; otherwise it is allocated on the - fly here — from the symm-mem pool in zero-copy mode (one-sided target), or a plain - tensor in normal mode (keeps allocation torch.compile / CUDA-graph safe and lets - autograd own the grad's lifetime). + """Autograd combine; bwd scatters the expert_out grad into ``grad_out``. When the caller + supplies it that buffer is used as-is; otherwise it is allocated in the backward from the + symm-mem pool in zero-copy mode (one-sided target) or a plain tensor in normal mode (keeps + allocation torch.compile / CUDA-graph safe and lets autograd own the grad's lifetime). - ``grad_out`` is the backward's scatter target (an output it writes, never reads), - so it is stashed as a plain ctx attribute rather than via save_for_backward, which - would version-track a tensor we mutate. + ``grad_out`` is a write-only scatter target, so it is stashed as a plain ctx attribute rather + than via save_for_backward, which would version-track a tensor we mutate. """ @staticmethod @@ -521,48 +609,89 @@ def forward( # type: ignore[override] hidden_dim: int, grad_out: Optional[torch.Tensor], expert_out: torch.Tensor, + bwd_quant_recipe=None, + token_counts: Optional[torch.Tensor] = None, ): - """Combine fwd; stashes the bwd grad target or expert_out shape to size it.""" + """Combine fwd; stashes the bwd grad target or expert_out shape to size it. When + ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8.""" device = expert_out.device result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) ctx.save_for_backward(handle_mem) ctx.grad_out = grad_out - if grad_out is None: - ctx.expert_out_shape = expert_out.shape - ctx.expert_out_dtype = expert_out.dtype - ctx.device = device + ctx.bwd_quant_recipe = bwd_quant_recipe + ctx.token_counts = token_counts + ctx.expert_out_shape = expert_out.shape + ctx.expert_out_dtype = expert_out.dtype + ctx.device = device return result @staticmethod def backward(ctx, g_result): # type: ignore[override] - """Combine bwd; scatters the result grad into the grad target.""" + """Combine bwd; scatters the result-grad to expert positions. High-precision sends the grad + as-is; a quantized recipe (MXFP8 today) quantizes it and returns the expert_out grad as a + per-expert GroupedTensor. New FP8/NVFP4 recipes plug in via ``_quantize``.""" if not g_result.is_contiguous(): g_result = g_result.contiguous() (handle_mem,) = ctx.saved_tensors - grad_expert_out = ctx.grad_out - if grad_expert_out is None: - grad_expert_out = _alloc_io( - ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() + + if ctx.bwd_quant_recipe is None: + grad_expert_out = ctx.grad_out + if grad_expert_out is None: + grad_expert_out = _alloc_io( + ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() + ) + torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + else: + if tex.ep_get_zero_copy(): + raise NotImplementedError( + "Quantized combine backward is not supported under zero-copy" + ) + mx, g_scale_inv = _quantize(g_result, ctx.bwd_quant_recipe) + g_data = mx._rowwise_data + recv_pr, hidden = ctx.expert_out_shape[0], ctx.expert_out_shape[-1] + ge_data, ge_scale_inv = _scale_alloc_io( + ctx.grad_out, + recv_pr, + hidden, + g_scale_inv.shape[-1], + g_data.dtype, + g_scale_inv.dtype, + ctx.device, + False, + ) + # The backend keys on the fp8 scaling mode; reinterpret the byte-backed data as fp8. + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + g_data.view(torch.float8_e4m3fn), + ge_data.view(torch.float8_e4m3fn), + g_scale_inv, + ge_scale_inv, + ) + grad_expert_out = _make_grouped_mxfp8( + ge_data, ge_scale_inv, ctx.token_counts, mx._fp8_dtype, ctx.expert_out_dtype ) - torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + return ( None, # handle_mem None, # num_local_tokens None, # hidden_dim None, # grad_out grad_expert_out, + None, # bwd_quant_recipe + None, # token_counts ) # Public high-level wrappers -# NCCL EP currently only supports bfloat16 payload tensors. +# NCCL EP inputs are bfloat16; MXFP8 is applied internally via the buffer's dispatch_fwd_quant_recipe. def _require_bf16(name: str, t: torch.Tensor) -> None: if t.dtype is not torch.bfloat16: raise NotImplementedError( - f"NCCL EP currently supports only bfloat16 payloads; got {name}.dtype={t.dtype}." + "NCCL EP currently supports only bfloat16 or MXFP8 payloads; got" + f" {name}.dtype={t.dtype}." ) @@ -578,6 +707,116 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso return torch.empty(*shape, dtype=dtype, device=device) +def _quantize(x: torch.Tensor, recipe): + """Quantize a high-precision tensor per a block-scaled ``recipe`` and return + ``(quantized_tensor, scale_inv)`` where ``scale_inv`` is the compact ``[T, H/block]`` scale the + EP backend routes. The quantized tensor is returned so callers can keep it as the autograd + operand; its ``_rowwise_data`` is the fp8 payload and ``scale_inv.shape[-1]`` the scale-column + count. Only ``MXFP8BlockScaling`` is supported today; add a recipe branch below for new + FP8/NVFP4 recipes.""" + from ..common.recipe import MXFP8BlockScaling + + if isinstance(recipe, MXFP8BlockScaling): + return _quantize_mxfp8(x) + raise NotImplementedError( + f"EP block-scaled dispatch supports MXFP8BlockScaling only; got {type(recipe).__name__}." + ) + + +def _quantize_mxfp8(x: torch.Tensor): + """MXFP8 implementation of ``_quantize``. EP routes and returns E4M3 data in both directions, + so quantize to E4M3 regardless of pass. Strips the GEMM scale row padding to the compact + ``[T, H/block]`` layout; requires a 16-byte-aligned scale row.""" + from .constants import MXFP8_BLOCK_SCALING_SIZE + from .tensor.mxfp8_tensor import MXFP8Quantizer + + mx = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False).quantize(x) + if mx._with_gemm_swizzled_scales: + raise RuntimeError( + "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact." + ) + data = mx._rowwise_data + scale_inv = mx._rowwise_scale_inv + if data is None or scale_inv is None: + raise ValueError("MXFP8 tokens must carry rowwise data and scale_inv for EP dispatch.") + t_flat = x.shape[0] + hidden = x.shape[-1] + cols = hidden // MXFP8_BLOCK_SCALING_SIZE + # The backend forwards each token's scale row with a 16-byte-aligned store, so the row + # (cols * dtype bytes) must be a multiple of 16. + scale_row_bytes = cols * scale_inv.element_size() + if scale_row_bytes % 16 != 0: + raise ValueError( + f"MXFP8 dispatch requires a 16-byte-aligned scale row; hidden={hidden} gives " + f"{scale_row_bytes} bytes. Use a hidden size that is a multiple of " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}." + ) + # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block] + # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the + # slice stays contiguous; assert rather than force a copy. + scale_inv = scale_inv[:t_flat, :cols] + if not scale_inv.is_contiguous(): + raise ValueError( + "MXFP8 dispatch requires compact contiguous scales [T, H/block]; got a " + f"non-contiguous [{t_flat}, {cols}] slice." + ) + return mx, scale_inv + + +def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, device, zero_copy): + """Block-scaled output data + scale buffers, each ``rows`` tall, laid out back-to-back + (``[rows, data_cols]`` data of ``data_dtype`` then ``[rows, scale_cols]`` scales of + ``scale_dtype``). Carve both from a single caller ``buf`` when it is large enough, so one + symm-mem window backs both views; else allocate them (symm-mem pool under zero-copy, else + plain). Recipe-agnostic: byte sizes come from the element sizes.""" + data_bytes = rows * data_cols * torch.empty((), dtype=data_dtype).element_size() + scale_bytes = rows * scale_cols * torch.empty((), dtype=scale_dtype).element_size() + if buf is not None: + # Reinterpret in place; a non-contiguous buf would force a copy and leave the caller's + # buffer unwritten, so require contiguous and view rather than reshape. + if not buf.is_contiguous(): + raise ValueError("scaled output buffer must be contiguous.") + flat = buf.view(-1).view(torch.uint8) + if flat.numel() < data_bytes + scale_bytes: + raise ValueError( + f"scaled output buffer too small: need {data_bytes + scale_bytes} bytes " + f"(data + scales), got {flat.numel()}." + ) + data = flat[:data_bytes].view(data_dtype).reshape(rows, data_cols) + scale_inv = ( + flat[data_bytes : data_bytes + scale_bytes].view(scale_dtype).reshape(rows, scale_cols) + ) + return data, scale_inv + data = _alloc_io((rows, data_cols), data_dtype, device, zero_copy) + scale_inv = _alloc_io((rows, scale_cols), scale_dtype, device, zero_copy) + return data, scale_inv + + +def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype): + """Wrap expert-major MXFP8 recv data + compact e8m0 scales as a per-expert ``GroupedTensor``. + + ``token_counts`` (int64 [num_local_experts]) is the padded per-expert row counts (128-aligned), + used as the group sizes. Grouping is device-side (first_dims/tensor_offsets), so the counts never + sync to host; the outer shape is the static recv capacity, bounded per expert by first_dims. + """ + from .tensor.grouped_tensor import GroupedTensor + from .tensor.mxfp8_tensor import MXFP8Quantizer + + assert data.dim() == 2, "recv data must be 2D [capacity_rows, hidden]" + capacity_rows, hidden = data.shape + quantizer = MXFP8Quantizer(fp8_dtype, rowwise=True, columnwise=False) + return GroupedTensor( + shape=(capacity_rows, hidden), + dtype=fake_dtype, + num_tensors=token_counts.numel(), + quantizer=quantizer, + data=data.reshape(-1).detach(), + scale_inv=scale_inv.reshape(-1).detach(), + first_dims=token_counts, + tensor_offsets=tex.splits_to_offsets(token_counts, hidden), + ) + + def ep_dispatch( buffer: EpBuffer, tokens: torch.Tensor, @@ -587,39 +826,58 @@ def ep_dispatch( recv_tokens: Optional[torch.Tensor] = None, recv_topk_weights: Optional[torch.Tensor] = None, ): - """Prepare + dispatch with autograd. topk_idx must be int32 or int64. - - ``recv_tokens`` / ``recv_topk_weights`` are the dispatch recv outputs: pass caller-owned buffers - (mcore-managed mode; in zero-copy they must be symm-mem-backed) or leave them None to allocate on - the fly (zero-copy: symm-mem pool; normal: plain). In eager mode the recv outputs are sized to - this step's recv-token total and must not be caller-supplied. Returns (recv_tokens, - recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See ``buffer.total_recv_tokens`` for - the per-step recv total. + """Prepare + dispatch with autograd. ``tokens`` is bfloat16; ``topk_idx`` is int32 or int64. + + When the buffer's ``dispatch_fwd_quant_recipe`` is set (``MXFP8BlockScaling`` only for now), tokens + are quantized internally and recv is returned as a per-expert ``GroupedTensor``; otherwise recv + stays bfloat16. A pre-quantized ``tokens`` is not accepted. + + ``recv_tokens`` / ``recv_topk_weights`` are the recv outputs: pass caller-owned buffers + (symm-mem-backed under zero-copy) or leave them None to allocate. For MXFP8 the recv data and + scales share ``recv_tokens`` (data then scales), so size it to at least + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes. Eager mode sizes the recv outputs + per step and forbids caller-supplied buffers. + + Returns (recv_tokens, recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See + ``buffer.total_recv_tokens`` for the per-step recv total. """ - _require_bf16("tokens", tokens) + from .quantized_tensor import QuantizedTensor + if topk_weights.dtype is not torch.float32: raise TypeError( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) + if isinstance(tokens, QuantizedTensor): + raise NotImplementedError( + "NCCL EP dispatch takes a bfloat16 input and quantizes internally when the buffer's " + "dispatch_fwd_quant_recipe is set; a pre-quantized tensor is not accepted." + ) + _require_bf16("tokens", tokens) if buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): raise ValueError( "eager mode sizes the recv outputs from the per-step recv-token total " "and cannot use caller-supplied recv_tokens / recv_topk_weights" ) + # Prepare (routing AllGather) up front so the recv outputs can be sized; in # eager mode ep_prepare also host-syncs this step's recv-token total. tokens_per_expert = ep_prepare(buffer, topk_idx) rows = buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank - if recv_tokens is None: - recv_tokens = _alloc_io( - (rows, buffer.hidden_dim), - buffer.payload_dtype, - buffer.device, - buffer.zero_copy, - ) - if recv_topk_weights is None: - recv_topk_weights = _alloc_io((rows,), torch.float32, buffer.device, buffer.zero_copy) + + tokens_scale_inv = None + if buffer.dispatch_fwd_quant_recipe is not None: + # The grouped recv's packed scales only match the grouped layout when each expert slot is + # 128-aligned. + if buffer.alignment <= 0 or buffer.alignment % 128 != 0: + raise ValueError( + "MXFP8 dispatch requires a per-expert alignment that is a positive multiple of " + f"128; got alignment={buffer.alignment}." + ) + # Quantize here (not in forward) so the quantized tensor stays the autograd operand and grad + # reaches the pre-quant input; forward then carves the recv buffers and routes. + tokens, tokens_scale_inv = _quantize(tokens, buffer.dispatch_fwd_quant_recipe) + recv_tokens, recv_topk_weights = _EpDispatch.apply( buffer.handle_mem, recv_tokens, @@ -627,6 +885,10 @@ def ep_dispatch( topk_idx, tokens, topk_weights, + tokens_scale_inv, + tokens_per_expert, + rows, + buffer.payload_dtype, ) return recv_tokens, recv_topk_weights, tokens_per_expert @@ -640,11 +902,12 @@ def ep_combine( ): """Combine with autograd; caller pre-applies topk weighting. - ``expert_out`` is the combine input (always caller-supplied; in zero-copy it must be symm-mem- - backed). ``grad_out`` is the backward's grad target: pass a caller-owned buffer (mcore-managed - mode) or leave it None to allocate on the fly in the backward (zero-copy: symm-mem pool; normal: - plain). In eager mode the grad target is sized per step and must not be caller-supplied. Result - shape is (num_local_tokens, buffer.hidden_dim); defaults to buffer.max_tokens_per_rank rows. + ``expert_out`` is the combine input (symm-mem-backed under zero-copy). ``grad_out`` is the + backward's grad target: pass a caller-owned buffer or leave it None to allocate. For MXFP8 the + grad data and scales share ``grad_out`` (data then scales), so size it to at least + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes (non-zero-copy only). Eager mode sizes + the grad target per step and forbids a caller-supplied buffer. Result shape is + (num_local_tokens, hidden_dim); num_local_tokens defaults to buffer.max_tokens_per_rank. """ _require_bf16("expert_out", expert_out) if buffer.eager and grad_out is not None: @@ -654,10 +917,31 @@ def ep_combine( ) if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank + # When combine_bwd_quant_recipe is set the combine backward sends the result-grad over the + # wire as MXFP8 and returns the expert_out grad as a GroupedTensor. + bwd_quant_recipe = None + if buffer.combine_bwd_quant_recipe is not None: + from ..common.recipe import MXFP8BlockScaling + + if not isinstance(buffer.combine_bwd_quant_recipe, MXFP8BlockScaling): + raise NotImplementedError( + "EP combine backward supports MXFP8BlockScaling only; got " + f"{type(buffer.combine_bwd_quant_recipe).__name__}." + ) + # The grouped grad's packed scales only match the grouped layout when each expert slot is + # 128-aligned. + if buffer.alignment <= 0 or buffer.alignment % 128 != 0: + raise ValueError( + "MXFP8 combine backward requires a per-expert alignment that is a positive" + f" multiple of 128; got alignment={buffer.alignment}." + ) + bwd_quant_recipe = buffer.combine_bwd_quant_recipe return _EpCombine.apply( buffer.handle_mem, num_local_tokens, buffer.hidden_dim, grad_out, expert_out, + bwd_quant_recipe, + buffer.tokens_per_expert, )