fix lint and comments.

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

Author: fanshiqing

tests/pytorch/distributed/test_gtp.py

@@ -1604,3 +1604,55 @@ def test_fallback_skipped_when_already_finalized(self):
         assert torch.all(
             p.main_grad == 5.0
         ), "main_grad must be untouched when _already_finalized=True"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+    def test_fallback_skipped_for_pure_ag_param(self):
+        """Regression: cross-graph fwd-AG prefetch in flight + finalize_after_drain=True.
+
+        A param can be in _inflight_comm_params because of an outstanding async
+        all-gather (e.g. a cross-graph forward prefetch reaching the
+        bwd→optimizer boundary).  No reduce-scatter was ever issued for that
+        param, so _rs_ticket is None on every weight.  Previously the fallback
+        called cache.get(None) and crashed with KeyError; the guard now skips
+        the inline accumulation entirely when no weight has an RS ticket.
+        """
+        dtype = torch.bfloat16
+        p = GTPShardedParam(torch.zeros(8, 4, dtype=dtype, device="cuda"))
+        p.group = self._FakeGroup()
+        p.expert_idx = None
+        p.pad_length = 0
+        p.chain_id = gtp_module.GTPChain.UNGRAPHED.value
+        p._quantizer = None
+        p.is_routed_expert = False
+        # Pre-existing main_grad with a sentinel that must survive untouched.
+        p.main_grad = torch.full((8, 4), 7.0, dtype=dtype, device="cuda")
+        p._prefetch_handle = None
+        p._wgrad_rs_handle = None
+        p._cached_ag_stream = None
+        p._cached_rs_stream = None
+        p.ag_event = torch.cuda.Event(external=True)
+        p.rs_event = torch.cuda.Event(external=True)
+        p.rs_event.record()
+        p._already_finalized = False
+        # Critical: simulates a pure-AG prefetch — no RS ever issued, ticket is None.
+        p._rs_ticket = None
+
+        saved = set(gtp_module._inflight_comm_params)
+        gtp_module._inflight_comm_params.clear()
+        gtp_module._inflight_comm_params.add(p)
+        try:
+            # Must NOT raise KeyError(None) from cache.get(None).
+            gtp_module.wait_async_comms(
+                chain_id=p.chain_id,
+                skip_rs=False,
+                finalize_after_drain=True,
+            )
+        finally:
+            gtp_module._inflight_comm_params.clear()
+            gtp_module._inflight_comm_params.update(saved)
+
+        torch.cuda.synchronize()
+        assert torch.all(p.main_grad == 7.0), \
+            "main_grad must be untouched for a pure-AG param (no wgrad to accumulate)"
+        assert p._already_finalized is False, \
+            "_already_finalized must stay False — no finalize happened for a pure-AG param"

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -133,10 +133,10 @@ void init_extension() {
 
 #include "common/util/pybind_helper.h"
 
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m)
-  m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"),
-        py::arg("output") = py::none(), py::arg("noop") = py::none());
+// Bindings for the NVFP4 split-amax fast-path used by GTP (coalesced
+// per-expert amax allreduce + cast).  Kept separate from PYBIND11_MODULE
+// so the latter stays under cpplint's 500-line readability/fn_size limit.
+static void RegisterNvfp4AmaxBindings(py::module &m) {
   m.def("compute_amax_nvfp4", transformer_engine::pytorch::compute_amax_nvfp4,
         "NVFP4: compute local amax into output's amax buffers; no cast, no allreduce",
         py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none());
@@ -148,6 +148,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("compute_multi_amax_nvfp4", transformer_engine::pytorch::compute_multi_amax_nvfp4,
         "NVFP4: fused multi-tensor amax compute (writes both rowwise+columnwise amax per output)",
         py::arg("tensor_list"), py::arg("quantizer_list"), py::arg("output_list"));
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m)
+  m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"),
+        py::arg("output") = py::none(), py::arg("noop") = py::none());
+  RegisterNvfp4AmaxBindings(m);
   m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"),
         py::arg("otype"));
   m.def("create_empty_quantized_tensor",

transformer_engine/pytorch/module/generalized_tensor_parallelism.py

@@ -1679,11 +1679,16 @@ def wait_async_comms(
             param._already_ag_drained = True
         if not skip_rs:
             param._wait_reduce_scatter(finalize_grad=finalize_after_drain)
-            if finalize_after_drain and not getattr(param, "_already_finalized", False):
-                # Fallback path: _wait_reduce_scatter ran sync (no handle to
-                # drain), so it didn't accumulate. Do the inline accumulation
-                # here, matching the wgrad_reduce_scatter / _wait_reduce_scatter
-                # finalize-grad pattern (state reset + main_grad.add_ + release).
+            # Fallback inline-accumulation: only when finalize is requested,
+            # _wait_reduce_scatter didn't already finalize, and an RS actually
+            # ran for this param (rs_ticket set).  Skips pure-AG prefetches in
+            # _inflight_comm_params (no wgrad to accumulate).
+            need_fallback_accumulation = (
+                finalize_after_drain
+                and not getattr(param, "_already_finalized", False)
+                and any(w._rs_ticket is not None for w in param._weights)
+            )
+            if need_fallback_accumulation:
                 cache = get_global_GTP_cache()
                 param.rs_event.wait()
                 for w in param._weights:

transformer_engine/pytorch/module/linear.py

@@ -154,7 +154,7 @@ class LinearFwdArgs:
     cpu_offloading: bool
     is_grad_enabled: bool
 
-    # --- Extended tensor parallelism ---
+    # --- Generalized tensor parallelism ---
     gtp_size: int = 1
 
 
@@ -226,7 +226,7 @@ class LinearBwdArgs:
     cpu_offloading: bool = False
     owns_input: bool = False
 
-    # --- Extended tensor parallelism ---
+    # --- Generalized tensor parallelism ---
     gtp_size: int = 1
 
     # --- Per-backward scratch state (populated inside _linear_backward) ---