[PyTorch] Address PR #3009 review: remove .view() calls, int routing_map_format

Apply the four CPU-overhead fixes the reviewer asked for and the
CLAUDE.md "CPU overhead in PyTorch wrappers" section codifies:

1. _validate_routing_map_format returns plain int (not enum); the
   autograd Function + tex.* bindings only see ints. Validates via
   precomputed frozenset and a single dict.get with canonical
   lowercase keys (no .lower()/.upper()).

2. Type annotations on Function.forward use int (not the string
   forward-ref 'RoutingMapFormat').

3. Removed every .view() from FusedTopkScoreFunction.{forward,backward}
   and FusedComputeScoresForMoEAuxLoss.{forward,backward}. C++
   extension now accepts N-D logits/grad_probs, computes num_tokens
   from the product of leading dims, num_experts from the last dim,
   allocates outputs at the user-facing N-D shape, and wraps tensors
   with an explicit 2D shape via makeTransformerEngineTensor only for
   the kernel call. Asserts is_contiguous() on inputs.

4. Bwd allocates grad_logits with torch.empty_like(grad_probs) (N-D)
   instead of allocate-2D-then-view.

PyTorch-extension boundary takes 'int routing_map_format' and casts
to NVTERoutingMapFormat inside; the common-layer C API (nvte_*_v2)
keeps the enum.

Signed-off-by: tdophung <tdophung@nvidia.com>

Author: tdophung

transformer_engine/pytorch/csrc/extensions.h

@@ -26,25 +26,29 @@ namespace transformer_engine::pytorch {
  * Router fusion
  **************************************************************************************************/
 
+// PyTorch-extension boundary uses int for routing_map_format (not the enum) to
+// keep the Python hot path free of pybind11 enum construction. The int is
+// cast to NVTERoutingMapFormat once inside each function. See CLAUDE.md
+// "CPU overhead in PyTorch wrappers".
 std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_topk_with_score_function_fwd(
     at::Tensor logits, int topk, bool use_pre_softmax, std::optional<int> num_groups,
     std::optional<int> group_topk, std::optional<float> scaling_factor, std::string score_function,
     std::optional<at::Tensor> expert_bias,
-    NVTERoutingMapFormat routing_map_format = NVTE_ROUTING_MAP_FORMAT_BYTEMAP);
+    int routing_map_format = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP));
 
 void fused_topk_with_score_function_bwd(
-    int num_tokens, int num_experts, at::Tensor routing_map, at::Tensor intermediate_output,
-    at::Tensor grad_probs, at::Tensor grad_logits, int topk, bool use_pre_softmax,
-    std::optional<float> scaling_factor, std::string score_function,
-    NVTERoutingMapFormat routing_map_format = NVTE_ROUTING_MAP_FORMAT_BYTEMAP);
+    at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs,
+    at::Tensor grad_logits, int topk, bool use_pre_softmax, std::optional<float> scaling_factor,
+    std::string score_function,
+    int routing_map_format = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP));
 
 std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_score_for_moe_aux_loss_fwd(
     at::Tensor logits, int topk, std::string score_function,
-    NVTERoutingMapFormat routing_map_format = NVTE_ROUTING_MAP_FORMAT_BYTEMAP);
+    int routing_map_format = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP));
 
-void fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts,
-                                      at::Tensor intermediate_output, at::Tensor grad_probs,
-                                      at::Tensor grad_logits, int topk, std::string score_function);
+void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor grad_scores,
+                                      at::Tensor grad_logits, int topk,
+                                      std::string score_function);
 
 std::tuple<at::Tensor, at::Tensor> fused_moe_aux_loss_fwd(at::Tensor probs,
                                                           at::Tensor tokens_per_expert,

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -136,24 +136,27 @@ void init_router_bindings(pybind11::module &m) {
   pybind11::enum_<NVTERoutingMapFormat>(m, "NVTERoutingMapFormat", pybind11::module_local())
       .value("BYTEMAP", NVTE_ROUTING_MAP_FORMAT_BYTEMAP)
       .value("BITMAP_U8", NVTE_ROUTING_MAP_FORMAT_BITMAP_U8);
+  // routing_map_format is passed as int (not the enum) on the PyTorch hot
+  // path; see CLAUDE.md "CPU overhead in PyTorch wrappers".
   m.def("fused_topk_with_score_function_fwd", &fused_topk_with_score_function_fwd,
         py::arg("logits"), py::arg("topk"), py::arg("use_pre_softmax"), py::arg("num_groups"),
         py::arg("group_topk"), py::arg("scaling_factor"), py::arg("score_function"),
-        py::arg("expert_bias"), py::arg("routing_map_format") = NVTE_ROUTING_MAP_FORMAT_BYTEMAP,
+        py::arg("expert_bias"),
+        py::arg("routing_map_format") = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
         "Fused topk with score function fwd");
   m.def("fused_topk_with_score_function_bwd", &fused_topk_with_score_function_bwd,
-        py::arg("num_tokens"), py::arg("num_experts"), py::arg("routing_map"),
-        py::arg("intermediate_output"), py::arg("grad_probs"), py::arg("grad_logits"),
-        py::arg("topk"), py::arg("use_pre_softmax"), py::arg("scaling_factor"),
-        py::arg("score_function"), py::arg("routing_map_format") = NVTE_ROUTING_MAP_FORMAT_BYTEMAP,
+        py::arg("routing_map"), py::arg("intermediate_output"), py::arg("grad_probs"),
+        py::arg("grad_logits"), py::arg("topk"), py::arg("use_pre_softmax"),
+        py::arg("scaling_factor"), py::arg("score_function"),
+        py::arg("routing_map_format") = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
         "Fused topk with score function bwd");
   m.def("fused_score_for_moe_aux_loss_fwd", &fused_score_for_moe_aux_loss_fwd, py::arg("logits"),
         py::arg("topk"), py::arg("score_function"),
-        py::arg("routing_map_format") = NVTE_ROUTING_MAP_FORMAT_BYTEMAP,
+        py::arg("routing_map_format") = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
         "Fused aux loss with score function fwd");
   m.def("fused_score_for_moe_aux_loss_bwd", &fused_score_for_moe_aux_loss_bwd,
-        py::arg("num_tokens"), py::arg("num_experts"), py::arg("intermediate_output"),
-        py::arg("grad_scores"), py::arg("grad_logits"), py::arg("topk"), py::arg("score_function"),
+        py::arg("intermediate_output"), py::arg("grad_scores"), py::arg("grad_logits"),
+        py::arg("topk"), py::arg("score_function"),
         "Fused aux loss with score function bwd");
   m.def("fused_moe_aux_loss_fwd", &fused_moe_aux_loss_fwd, py::arg("probs"),
         py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"),