[Common][PyTorch] Add NVFP4 EP dispatch and combine backward

transformer_engine/common/CMakeLists.txt

@@ -364,6 +364,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources
      cast/cast_grouped.cu
      cast/cast_grouped_dbias.cu
      cast/mxfp8/requantize_mxfp8.cu
+     ep/nvfp4_ep_rows.cu
      fused_rope/fused_mla_kv_rope_mxfp8.cu
      gemm/cutlass_grouped_gemm.cu
      hadamard_transform/group_hadamard_transform.cu

transformer_engine/common/include/transformer_engine/ep.h

@@ -309,6 +309,75 @@ void nvte_ep_combine_bwd_v2(NVTETensor handle_mem, NVTETensor grad, NVTECommWind
                             NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win,
                             const NVTEEpLayerConfig* layer_cfg, cudaStream_t stream);
 
+/* -- NVFP4 token rows ------------------------------------------------------ */
+/* Pack, decode and receive kernels for moving NVFP4 tokens through the
+ * unscaled EP dispatch / combine-backward path. They do not need NCCL EP and
+ * are available in every build. Pack and receive need Blackwell (SM 10.0+).
+ *
+ * Packed row layout, row_bytes per token (a multiple of 16 and at least
+ * H/2 + H/16 + 4): [H/2 bytes of E2M1 pairs (low nibble first) | H/16 E4M3
+ * block scales (one per 16 elements) | FP32 token decode scale | zero pad].
+ * The token decode scale is a power of two: the smallest 2^e with
+ * amax <= 6 * 448 * 2^e (0 for an all-zero token). */
+
+/*! \brief Quantize BF16 tokens to packed NVFP4 rows.
+ *
+ *  \param[in]     input          BF16 tokens [T, H]; H % 512 == 0 and H <= 16384.
+ *  \param[out]    packed         Packed rows, kByte [T, row_bytes].
+ *  \param[in]     rng_state      Optional (may be NULL) int64 [2] (seed, offset) on device.
+ *                                When given, E2M1 values are stochastically rounded; else RN.
+ *  \param[in]     use_fast_math  Use the approximate reciprocal, as NVTE_USE_FAST_MATH does for
+ *                                TE's NVFP4 casts.
+ *  \param[in]     stream         CUDA stream.
+ */
+void nvte_ep_nvfp4_pack(const NVTETensor input, NVTETensor packed, const NVTETensor rng_state,
+                        int use_fast_math, cudaStream_t stream);
+
+/*! \brief Decode packed NVFP4 rows to BF16.
+ *
+ *  \param[in]     packed         Packed rows, kByte [R, row_bytes].
+ *  \param[in]     counts         Optional (may be NULL) int64 [G] on device. When given, only the
+ *                                first min(sum(counts), R) rows are decoded and the rest of
+ *                                ``output`` is left untouched (CUDA-graph safe).
+ *  \param[out]    output         BF16 [R, H]; H % 64 == 0.
+ *  \param[in]     stream         CUDA stream.
+ */
+void nvte_ep_nvfp4_unpack(const NVTETensor packed, const NVTETensor counts, NVTETensor output,
+                          cudaStream_t stream);
+
+/*! \brief Build the per-expert NVFP4 grouped-GEMM operands straight from packed rows.
+ *
+ *  Rows are grouped by expert: group g owns counts[g] consecutive rows, each count a multiple
+ *  of 128. The rowwise operand reuses the sender's E2M1 codes; its block scales are re-expressed
+ *  against the group scale G = max token decode scale of the group (exact while they stay in
+ *  E4M3's normal range). The columnwise operand is the 16-point random Hadamard transform along
+ *  tokens, quantized in the same pass. Second-level amaxes are row_amax[g] = 6 * 448 * G and
+ *  col_amax[g] = 4 * 6 * 448 * G (a bound of the post-RHT amax, so no amax pass is needed).
+ *  Scales use the GEMM-swizzled layout. Bounded on device by ``counts`` (CUDA-graph safe).
+ *
+ *  \param[in]     packed         Packed rows (see above), kByte [R, row_bytes]; R % 128 == 0.
+ *  \param[in]     counts         int64 [G] on device, 1 <= G <= 128.
+ *  \param[in,out] workspace      int32 [G] scratch.
+ *  \param[in]     sign_mask      RHT random signs: bit i set when sign i is -1.
+ *  \param[out]    row_data       kByte [R * H / 2].
+ *  \param[out]    row_scale      kByte [R * H / 16].
+ *  \param[out]    col_data       kByte [R * H / 2]; group g is [H, counts[g]].
+ *  \param[out]    col_scale      kByte [R * H / 16].
+ *  \param[out]    row_amax       float32 [G].
+ *  \param[out]    col_amax       float32 [G].
+ *  \param[in]     rng_state      Optional (may be NULL) int64 [2] (seed, offset) on device:
+ *                                stochastic rounding of the columnwise operand when given.
+ *  \param[in]     hidden         H; H % 256 == 0.
+ *  \param[in]     use_fast_math  Use the approximate reciprocal.
+ *  \param[in]     stream         CUDA stream.
+ */
+void nvte_ep_nvfp4_direct_receive(const NVTETensor packed, const NVTETensor counts,
+                                  NVTETensor workspace, uint32_t sign_mask, NVTETensor row_data,
+                                  NVTETensor row_scale, NVTETensor col_data, NVTETensor col_scale,
+                                  NVTETensor row_amax, NVTETensor col_amax,
+                                  const NVTETensor rng_state, size_t hidden, int use_fast_math,
+                                  cudaStream_t stream);
+
 #ifdef __cplusplus
 }
 #endif

transformer_engine/pytorch/csrc/extensions.h

@@ -824,6 +824,17 @@ void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expe
 // Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP.
 void register_ep_bindings(pybind11::module_ &m);
 
+// NVFP4 token rows for the EP dispatch / combine backward (independent of NCCL EP).
+at::Tensor ep_nvfp4_pack(const at::Tensor &tokens, bool stochastic);
+
+at::Tensor ep_nvfp4_unpack(const at::Tensor &packed, int64_t hidden,
+                           std::optional<at::Tensor> counts, std::optional<at::Tensor> out);
+
+// Returns {row_data, row_scale, col_data, col_scale, row_amax, col_amax}.
+std::vector<at::Tensor> ep_nvfp4_direct_receive(const at::Tensor &packed, int64_t hidden,
+                                                const at::Tensor &counts, int64_t sign_mask,
+                                                bool stochastic);
+
 /***************************************************************************************************
  * NVSHMEM APIs
  **************************************************************************************************/

transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp

@@ -0,0 +1,136 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+// NVFP4 token rows for the EP dispatch and combine backward. Independent of NCCL EP.
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <optional>
+#include <vector>
+
+#include "../common.h"
+#include "../extensions.h"
+#include "../pybind.h"
+#include "common/util/system.h"
+#include "transformer_engine/ep.h"
+
+namespace transformer_engine::pytorch {
+
+namespace {
+
+int64_t nvfp4_row_bytes(int64_t hidden) {
+  return static_cast<int64_t>(roundup(static_cast<size_t>(hidden / 2 + hidden / 16 + 4), 16));
+}
+
+int use_fast_math() { return transformer_engine::getenv<bool>("NVTE_USE_FAST_MATH") ? 1 : 0; }
+
+// (seed, offset) from torch's CUDA generator; under graph capture the offset is read on device,
+// so every replay draws fresh bits.
+at::Tensor make_rng_state(size_t elts_per_thread) {
+  auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
+      std::nullopt, at::cuda::detail::getDefaultCUDAGenerator());
+  at::PhiloxCudaState philox_args = init_philox_state(gen, elts_per_thread);
+  auto rng_state = at::empty({2}, at::TensorOptions().dtype(at::kLong).device(at::kCUDA));
+  philox_unpack(philox_args, static_cast<int64_t *>(rng_state.data_ptr()));
+  return rng_state;
+}
+
+}  // namespace
+
+at::Tensor ep_nvfp4_pack(const at::Tensor &tokens, bool stochastic) {
+  NVTE_CHECK(tokens.is_cuda() && tokens.dim() == 2, "tokens must be a 2D CUDA tensor.");
+  // Launch, allocate and draw RNG state on the input's device, not the current one.
+  const at::cuda::CUDAGuard device_guard(tokens.device());
+  const auto x = tokens.contiguous();
+  auto packed = at::empty({x.size(0), nvfp4_row_bytes(x.size(1))}, x.options().dtype(at::kByte));
+  std::optional<at::Tensor> rng_state;
+  TensorWrapper rng_state_nvte;
+  if (stochastic) {
+    rng_state = make_rng_state(4);
+    rng_state_nvte = makeTransformerEngineTensor(*rng_state);
+  }
+  auto x_nvte = makeTransformerEngineTensor(x);
+  auto packed_nvte = makeTransformerEngineTensor(packed);
+  NVTE_SCOPED_GIL_RELEASE({
+    nvte_ep_nvfp4_pack(x_nvte.data(), packed_nvte.data(), rng_state_nvte.data(), use_fast_math(),
+                       at::cuda::getCurrentCUDAStream());
+  });
+  return packed;
+}
+
+at::Tensor ep_nvfp4_unpack(const at::Tensor &packed, int64_t hidden,
+                           std::optional<at::Tensor> counts, std::optional<at::Tensor> out) {
+  NVTE_CHECK(packed.is_cuda() && packed.dim() == 2 && packed.scalar_type() == at::kByte,
+             "packed must be a 2D uint8 CUDA tensor.");
+  const at::cuda::CUDAGuard device_guard(packed.device());
+  const auto p = packed.contiguous();
+  if (!out) {
+    out = at::empty({p.size(0), hidden}, p.options().dtype(at::kBFloat16));
+  }
+  NVTE_CHECK(out->is_contiguous() && out->numel() == p.size(0) * hidden,
+             "out must be a contiguous [", p.size(0), ", ", hidden, "] tensor.");
+  auto out_2d = out->view({p.size(0), hidden});
+  TensorWrapper counts_nvte;
+  if (counts) {
+    counts = counts->contiguous();
+    counts_nvte = makeTransformerEngineTensor(*counts);
+  }
+  auto p_nvte = makeTransformerEngineTensor(p);
+  auto out_nvte = makeTransformerEngineTensor(out_2d);
+  NVTE_SCOPED_GIL_RELEASE({
+    nvte_ep_nvfp4_unpack(p_nvte.data(), counts_nvte.data(), out_nvte.data(),
+                         at::cuda::getCurrentCUDAStream());
+  });
+  return *out;
+}
+
+std::vector<at::Tensor> ep_nvfp4_direct_receive(const at::Tensor &packed, int64_t hidden,
+                                                const at::Tensor &counts, int64_t sign_mask,
+                                                bool stochastic) {
+  NVTE_CHECK(packed.is_cuda() && packed.dim() == 2 && packed.scalar_type() == at::kByte,
+             "packed must be a 2D uint8 CUDA tensor.");
+  NVTE_CHECK(counts.scalar_type() == at::kLong, "counts must be int64.");
+  const at::cuda::CUDAGuard device_guard(packed.device());
+  const auto p = packed.contiguous();
+  const auto c = counts.contiguous();
+  const int64_t capacity = p.size(0);
+  const int64_t num_groups = c.numel();
+  const auto u8 = p.options().dtype(at::kByte);
+  const auto f32 = p.options().dtype(at::kFloat);
+  auto row_data = at::empty({capacity * hidden / 2}, u8);
+  auto row_scale = at::empty({capacity * hidden / 16}, u8);
+  auto col_data = at::empty({capacity * hidden / 2}, u8);
+  auto col_scale = at::empty({capacity * hidden / 16}, u8);
+  auto row_amax = at::empty({num_groups}, f32);
+  auto col_amax = at::empty({num_groups}, f32);
+  auto workspace = at::empty({num_groups}, p.options().dtype(at::kInt));
+  std::optional<at::Tensor> rng_state;
+  TensorWrapper rng_state_nvte;
+  if (stochastic) {
+    rng_state = make_rng_state(32);
+    rng_state_nvte = makeTransformerEngineTensor(*rng_state);
+  }
+  auto p_nvte = makeTransformerEngineTensor(p);
+  auto c_nvte = makeTransformerEngineTensor(c);
+  auto workspace_nvte = makeTransformerEngineTensor(workspace);
+  auto row_data_nvte = makeTransformerEngineTensor(row_data);
+  auto row_scale_nvte = makeTransformerEngineTensor(row_scale);
+  auto col_data_nvte = makeTransformerEngineTensor(col_data);
+  auto col_scale_nvte = makeTransformerEngineTensor(col_scale);
+  auto row_amax_nvte = makeTransformerEngineTensor(row_amax);
+  auto col_amax_nvte = makeTransformerEngineTensor(col_amax);
+  NVTE_SCOPED_GIL_RELEASE({
+    nvte_ep_nvfp4_direct_receive(
+        p_nvte.data(), c_nvte.data(), workspace_nvte.data(), static_cast<uint32_t>(sign_mask),
+        row_data_nvte.data(), row_scale_nvte.data(), col_data_nvte.data(), col_scale_nvte.data(),
+        row_amax_nvte.data(), col_amax_nvte.data(), rng_state_nvte.data(),
+        static_cast<size_t>(hidden), use_fast_math(), at::cuda::getCurrentCUDAStream());
+  });
+  return {row_data, row_scale, col_data, col_scale, row_amax, col_amax};
+}
+
+}  // namespace transformer_engine::pytorch

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -372,6 +372,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 #ifdef NVTE_WITH_NCCL_EP
   transformer_engine::pytorch::register_ep_bindings(m);
 #endif  // NVTE_WITH_NCCL_EP
+  m.def("ep_nvfp4_pack", &transformer_engine::pytorch::ep_nvfp4_pack,
+        "Pack BF16 tokens into NVFP4 EP rows (power-of-two token scales)", py::arg("tokens"),
+        py::arg("stochastic"));
+  m.def("ep_nvfp4_unpack", &transformer_engine::pytorch::ep_nvfp4_unpack,
+        "Decode NVFP4 EP rows to BF16", py::arg("packed"), py::arg("hidden"),
+        py::arg("counts") = std::nullopt, py::arg("out") = std::nullopt);
+  m.def("ep_nvfp4_direct_receive", &transformer_engine::pytorch::ep_nvfp4_direct_receive,
+        "Build per-expert NVFP4 GEMM operands from NVFP4 EP rows", py::arg("packed"),
+        py::arg("hidden"), py::arg("counts"), py::arg("sign_mask"), py::arg("stochastic"));
 
   // Permutation functions
   m.def("moe_permute_fwd", transformer_engine::pytorch::moe_permute_fwd, "MOE permute FWD",