[SM100] NVFP4 MegaMoE: take the L2 activation scale as a device tensor

AI/nvfp4-mega-moe.md

@@ -80,16 +80,18 @@ deep_gemm.nvfp4_mega_moe(
     shared_l2_weights=shared_w2,
     activation_clamp=7.0, activation_alpha=1.702, activation_beta=1.0,
     l1_alpha=l1_global_products, l2_alpha=l2_global_products,
-    l2_activation_scale=mid_global_scale,
+    l2_activation_scale=mid_global_scale,  # shape (1,) float32 CUDA tensor
 )
 ```
 
 NVFP4 dequantization is `E2M1_value * E4M3_block_scale * global_scale`.
 The optional `l1_alpha` and `l2_alpha` are contiguous CUDA float32 vectors,
 one entry per local expert, containing activation-global × weight-global
 dequantization scales for each GEMM. They multiply FP32 accumulators before
-BF16 rounding. `l2_activation_scale` controls the fused intermediate quantizer
-and must also be included in `l2_alpha`. All default to one. Input global
+BF16 rounding. `l2_activation_scale`, a one-element CUDA float32 tensor, controls
+the fused intermediate quantizer and must also be included in `l2_alpha`. All
+default to one. The kernel reads all three at launch, so a CUDA graph captured
+once sees later in-place updates. Input global
 scaling must be consistent across ranks when using per-expert L1 products.
 
 `per_token_cast_to_nvfp4` and `cast_back_from_nvfp4` are PyTorch reference

csrc/apis/nvfp4_mega_moe.hpp

@@ -283,7 +283,7 @@ static void nvfp4_mega_moe(
     const float& activation_beta,
     const std::optional<torch::stable::Tensor>& l1_alpha_opt,
     const std::optional<torch::stable::Tensor>& l2_alpha_opt,
-    const float& l2_activation_scale,
+    const std::optional<torch::stable::Tensor>& l2_activation_scale,
     const bool& use_x_scales,
     const std::string& activation
 ) {
@@ -292,7 +292,6 @@ static void nvfp4_mega_moe(
 
     // Config checks
     const auto num_tokens = static_cast<int>(y.size(0));
-    DG_HOST_ASSERT(std::isfinite(l2_activation_scale) and l2_activation_scale > 0);
     DG_HOST_ASSERT(shared_l1_weights_tuple_opt.has_value() == shared_l2_weights_tuple_opt.has_value());
 
     // Activation checks
@@ -334,6 +333,11 @@ static void nvfp4_mega_moe(
             DG_HOST_ASSERT(alpha->device() == y.device());
         }
     }
+    if (l2_activation_scale.has_value()) {
+        DG_HOST_ASSERT(l2_activation_scale->scalar_type() == torch::headeronly::ScalarType::Float);
+        DG_HOST_ASSERT(l2_activation_scale->numel() == 1);
+        DG_HOST_ASSERT(l2_activation_scale->device() == y.device());
+    }
 
     // Packed E4M3 block-scale bytes, MN-major and TMA aligned.
     constexpr int kGranMN = 1;
@@ -484,7 +488,8 @@ static void nvfp4_mega_moe(
     double activation_alpha, double activation_beta,
     const std::optional<torch::stable::Tensor>& l1_alpha,
     const std::optional<torch::stable::Tensor>& l2_alpha,
-    double l2_activation_scale, bool use_x_scales, const std::string& activation) {
+    const std::optional<torch::stable::Tensor>& l2_activation_scale,
+    bool use_x_scales, const std::string& activation) {
     DG_HOST_ASSERT(shared_l1_weights.has_value() == shared_l2_weights.has_value());
     std::optional<std::tuple<torch::stable::Tensor, std::optional<torch::stable::Tensor>>> shared_l1;
     std::optional<std::tuple<torch::stable::Tensor, std::optional<torch::stable::Tensor>>> shared_l2;
@@ -504,7 +509,7 @@ static void nvfp4_mega_moe(
             : std::nullopt,
         fast_math, static_cast<float>(activation_alpha),
         static_cast<float>(activation_beta), l1_alpha, l2_alpha,
-        static_cast<float>(l2_activation_scale), use_x_scales, activation);
+        l2_activation_scale, use_x_scales, activation);
 }
 
 } // namespace deep_gemm::torch_registration
@@ -514,7 +519,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
     m.def("get_block_m_for_nvfp4_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_tokens, int num_topk) -> int");
     m.def("get_symm_buffer_size_for_nvfp4_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_topk, int hidden, int intermediate_hidden, int num_shared_experts=0, bool shared_bf16=False) -> (int, int[])");
     m.def("_slice_symm_buffer_for_nvfp4_mega_moe(Tensor(a) buffer, int[] layout_info) -> (Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a))");
-    m.def("nvfp4_mega_moe(Tensor(a!) y, Tensor l1_weights_tuple, Tensor l1_weights_tuple_sf, Tensor l2_weights_tuple, Tensor l2_weights_tuple_sf, Tensor? shared_l1_weights_tuple_opt, Tensor? shared_l1_weights_tuple_opt_sf, Tensor? shared_l2_weights_tuple_opt, Tensor? shared_l2_weights_tuple_opt_sf, Tensor(b!)? cumulative_local_expert_recv_stats, Tensor(c!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, float? activation_clamp_opt, bool fast_math, float activation_alpha, float activation_beta, Tensor? l1_alpha=None, Tensor? l2_alpha=None, float l2_activation_scale=1.0, bool use_x_scales=False, str activation='swiglu') -> ()");
+    m.def("nvfp4_mega_moe(Tensor(a!) y, Tensor l1_weights_tuple, Tensor l1_weights_tuple_sf, Tensor l2_weights_tuple, Tensor l2_weights_tuple_sf, Tensor? shared_l1_weights_tuple_opt, Tensor? shared_l1_weights_tuple_opt_sf, Tensor? shared_l2_weights_tuple_opt, Tensor? shared_l2_weights_tuple_opt_sf, Tensor(b!)? cumulative_local_expert_recv_stats, Tensor(c!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, float? activation_clamp_opt, bool fast_math, float activation_alpha, float activation_beta, Tensor? l1_alpha=None, Tensor? l2_alpha=None, Tensor? l2_activation_scale=None, bool use_x_scales=False, str activation='swiglu') -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(deep_gemm, CompositeExplicitAutograd, m) {

csrc/jit_kernels/impls/sm100_nvfp4_mega_moe.hpp

@@ -40,7 +40,7 @@ static void sm100_nvfp4_mega_moe(
     const bool& fast_math,
     const std::optional<torch::stable::Tensor>& l1_alpha,
     const std::optional<torch::stable::Tensor>& l2_alpha,
-    const float& l2_activation_scale,
+    const std::optional<torch::stable::Tensor>& l2_activation_scale,
     const bool& use_x_scales
 ) {
     const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
@@ -241,7 +241,7 @@ static void __instantiate_kernel() {{
         cumulative_local_expert_recv_stats_ptr,
         l1_alpha.has_value() ? l1_alpha->const_data_ptr<float>() : nullptr,
         l2_alpha.has_value() ? l2_alpha->const_data_ptr<float>() : nullptr,
-        l2_activation_scale,
+        l2_activation_scale.has_value() ? l2_activation_scale->const_data_ptr<float>() : nullptr,
         num_tokens,
         layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),
         tensor_map_l1_acts,

deep_gemm/_C.py

@@ -461,7 +461,7 @@ def nvfp4_mega_moe(
     activation_beta,
     l1_alpha=None,
     l2_alpha=None,
-    l2_activation_scale=1.0,
+    l2_activation_scale=None,
     use_x_scales=False,
     activation="swiglu",
 ):

deep_gemm/include/deep_gemm/impls/sm100_nvfp4_mega_moe.cuh

@@ -69,7 +69,7 @@ sm100_nvfp4_mega_moe_impl(void* y,
                             int* cumulative_local_expert_recv_stats,
                             const float* l1_alpha,
                             const float* l2_alpha,
-                            const float l2_activation_scale,
+                            const float* l2_activation_scale_ptr,
                             const uint32_t num_tokens,
                             const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
                             const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
@@ -1191,6 +1191,8 @@ sm100_nvfp4_mega_moe_impl(void* y,
         const auto epilogue_wg_idx = epilogue_warp_idx / 4;
         const auto epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx;
         const auto warp_idx_in_wg = epilogue_warp_idx % 4;
+        // Read at launch, so a CUDA graph replay sees the current value
+        const float l2_activation_scale = l2_activation_scale_ptr == nullptr ? 1.0f : *l2_activation_scale_ptr;
         DG_STATIC_ASSERT((kNumDispatchWarps + kNumMMANonEpilogueWarps) % 4 == 0 and
                          kNumEpilogueWarps % 4 == 0, "Invalid epilogue warps");
 

deep_gemm/mega/nvfp4.py

@@ -92,7 +92,7 @@ def nvfp4_mega_moe(y: torch.Tensor,
                   activation_beta: float = 0.0,
                   l1_alpha: Optional[torch.Tensor] = None,
                   l2_alpha: Optional[torch.Tensor] = None,
-                  l2_activation_scale: float = 1.0,
+                  l2_activation_scale: Optional[torch.Tensor] = None,
                   use_x_scales: bool = False,
                   activation: str = 'swiglu'):
     """Fused NVFP4 routed experts and MXFP8 or BF16 shared experts (SM100/SM103).
@@ -107,8 +107,10 @@ def nvfp4_mega_moe(y: torch.Tensor,
     Optional CUDA float32 vectors ``l1_alpha`` and ``l2_alpha`` (one entry per
     local expert) multiply GEMM accumulators before BF16 rounding. They contain
     the products of activation and weight global dequantization scales.
-    ``l2_activation_scale`` is the positive global dequantization scale used
-    when quantizing routed SwiGLU outputs; include it in ``l2_alpha``.
+    ``l2_activation_scale`` is an optional one-element CUDA float32 tensor holding
+    the positive global dequantization scale used when quantizing routed SwiGLU
+    outputs; include it in ``l2_alpha``. Like ``l1_alpha`` and ``l2_alpha``, it is
+    read at launch, so a CUDA graph replay sees values written after capture.
     ``activation='swiglustep'`` clips the gate after SiLU and supports routed
     experts only, with alpha=1 and beta=0. The default clips before SiLU.
     Defaults use unit global scales. The routed L2 input is the unweighted
@@ -125,6 +127,8 @@ def nvfp4_mega_moe(y: torch.Tensor,
         assert activation_alpha == 1.0 and activation_beta == 0.0
         assert sym_buffer.num_shared_experts == 0, 'swiglustep supports routed experts only'
     assert isinstance(sym_buffer, NVFP4SymmBuffer), 'NVFP4 requires NVFP4SymmBuffer'
+    assert l2_activation_scale is None or isinstance(l2_activation_scale, torch.Tensor), \
+        'l2_activation_scale must be a one-element CUDA float32 tensor'
     assert (shared_l1_weights is not None) == (sym_buffer.num_shared_experts > 0)
     assert (shared_l2_weights is not None) == (sym_buffer.num_shared_experts > 0)
     if isinstance(shared_l1_weights, torch.Tensor):

tests/test_nvfp4_mega_moe.py

@@ -41,6 +41,26 @@ def shared_sf_layout(sf, block_m, dst):
     dst[indices] = sf
 
 
+def check_graph_replay_reads_scales(kernel, y, transformed, buf, kwargs):
+    """A CUDA graph replay uses the L2 scales written after capture."""
+    kwargs = dict(kwargs, cumulative_local_expert_recv_stats=None)
+    scale, l2_alpha = kwargs['l2_activation_scale'], kwargs['l2_alpha']
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        kernel(y, *transformed, buf, **kwargs)
+    # Scale both together: a replay that kept the captured scale would double the output
+    scale.mul_(2)
+    l2_alpha.mul_(2)
+    graph.replay()
+    torch.cuda.synchronize()
+    replayed = y.clone()
+    kernel(y, *transformed, buf, **dict(kwargs, l2_activation_scale=scale.clone(), l2_alpha=l2_alpha.clone()))
+    torch.cuda.synchronize()
+    torch.testing.assert_close(replayed, y, atol=0.03125, rtol=0.02)
+    scale.div_(2)
+    l2_alpha.div_(2)
+
+
 def quantize(x, mode, scale=1.0):
     if mode == 'nvfp4':
         return per_token_cast_to_nvfp4(x, scale)
@@ -227,7 +247,8 @@ def gpu_process_snapshot():
                       activation=args.activation, activation_clamp=args.clamp, activation_alpha=args.alpha, activation_beta=args.beta,
                       fast_math=not args.exact_math)
         if nv:
-            kwargs.update(l1_alpha=l1_alpha, l2_alpha=l2_alpha, l2_activation_scale=args.mid_scale,
+            mid_scale = torch.full((1,), args.mid_scale, dtype=torch.float32, device='cuda')
+            kwargs.update(l1_alpha=l1_alpha, l2_alpha=l2_alpha, l2_activation_scale=mid_scale,
                           use_x_scales=args.x_scales)
         kernel = deep_gemm.nvfp4_mega_moe if nv else deep_gemm.fp8_fp4_mega_moe
         fn = lambda: kernel(y, *transformed, buf, **kwargs)
@@ -248,7 +269,9 @@ def gpu_process_snapshot():
                     lambda: kernel(y[:, :-1], *transformed, buf, **kwargs),
                     lambda: deep_gemm.fp8_fp4_mega_moe(y, *transformed, buf),
                     lambda: kernel(y, *transformed, buf, **dict(kwargs, l1_alpha=l1_alpha.double())),
-                    lambda: kernel(y, *transformed, buf, **dict(kwargs, l2_activation_scale=0.0)),
+                    lambda: kernel(y, *transformed, buf, **dict(kwargs, l2_activation_scale=args.mid_scale)),
+                    lambda: kernel(y, *transformed, buf, **dict(kwargs, l2_activation_scale=mid_scale.double())),
+                    lambda: kernel(y, *transformed, buf, **dict(kwargs, l2_activation_scale=mid_scale.repeat(2))),
                 ]
                 # Prove that Python dispatches to its own native entry point,
                 # without a hidden dependency on the legacy FP8/FP4 binding.
@@ -295,9 +318,11 @@ def gpu_process_snapshot():
             assert torch.equal(stats, count.chunk(world)[rank] * 2)
             torch.testing.assert_close(y, first, atol=0.03125, rtol=0.02)
             if nv and args.input_scale == args.weight_scale == args.mid_scale == 1.0:
-                kernel(y, *transformed, buf, **dict(kwargs, l1_alpha=None, l2_alpha=None,
+                kernel(y, *transformed, buf, **dict(kwargs, l1_alpha=None, l2_alpha=None, l2_activation_scale=None,
                                                    cumulative_local_expert_recv_stats=None))
                 torch.testing.assert_close(y, first, atol=0.03125, rtol=0.02)
+            if nv and m:
+                check_graph_replay_reads_scales(kernel, y, transformed, buf, kwargs)
                 assert torch.equal(stats, count.chunk(world)[rank] * 2)
         if args.benchmark:
             row.update(measure(fn, not args.warm_cache, args.repeats))