Make batched MoE GEMM CUDA-graph-safe with device-side alpha

- Persist Gemm object in MoeGemmState (avoids stack-local params_ destruction)
- Move gemm.initialize() to _init (triggers cudaFuncSetAttribute once)
- _run rebuilds params from arguments then calls gemm.run() (graph-safe)
- Change alpha from host float to device pointer (const float*); CUTLASS
  epilogue reads via alpha_ptr for zero host-GPU sync
- Update op definition, registered kernel, functional API, and benchmark

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

benchmarks/bench_moe_gemm_sm100.py

@@ -121,12 +121,14 @@ def bench_batched_nvfp4(lib, max_M, N, K, num_experts):
 
     lib.cgemm_nvfp4_moe_sm100_run.restype = ct.c_int
 
+    alpha_dev = torch.tensor([1.0], dtype=torch.float32, device=device)
+
     def run_kernel():
         lib.cgemm_nvfp4_moe_sm100_run(
             get_ptr(A_batched), get_ptr(B_batched),
             get_ptr(SFA), get_ptr(SFB),
             get_ptr(D_out),
-            ct.c_float(1.0), stream_ptr)
+            get_ptr(alpha_dev), stream_ptr)
 
     # Warmup
     for _ in range(WARMUP):

bitsandbytes/_ops.py

@@ -651,7 +651,7 @@ def _(
 torch.library.define(
     "bitsandbytes::gemm_nvfp4_moe",
     "(Tensor A_batched, Tensor B_batched, Tensor SFA, Tensor SFB, "
-    "float alpha, int max_M, int N, int K, int num_experts) -> Tensor",
+    "Tensor alpha, int max_M, int N, int K, int num_experts) -> Tensor",
 )
 
 
@@ -661,7 +661,7 @@ def _(
     B_batched: torch.Tensor,
     SFA: torch.Tensor,
     SFB: torch.Tensor,
-    alpha: float,
+    alpha: torch.Tensor,
     max_M: int,
     N: int,
     K: int,

bitsandbytes/backends/cuda/ops.py

@@ -1352,7 +1352,7 @@ def _(
     B_batched: torch.Tensor,
     SFA: torch.Tensor,
     SFB: torch.Tensor,
-    alpha: float,
+    alpha: torch.Tensor,
     max_M: int,
     N: int,
     K: int,
@@ -1377,13 +1377,16 @@ def _(
 
         _moe_batched_cache = {"key": key, "workspace": workspace}
 
+    # Ensure alpha is a float32 device tensor
+    alpha_dev = alpha.to(dtype=torch.float32, device=A_batched.device).contiguous()
+
     D_out = torch.empty(num_experts * max_M * N, dtype=torch.bfloat16, device=A_batched.device)
 
     ret = lib.cgemm_nvfp4_moe_sm100_run(
         get_ptr(A_batched), get_ptr(B_batched),
         get_ptr(SFA), get_ptr(SFB),
         get_ptr(D_out),
-        ct.c_float(alpha),
+        get_ptr(alpha_dev),
         ct.c_void_p(_get_tensor_stream(A_batched)),
     )
     if ret != 0:

bitsandbytes/functional.py

@@ -1543,10 +1543,9 @@ def gemm_nvfp4_grouped(
 def gemm_nvfp4_moe(
     A_batched: torch.Tensor,
     SFA_batched: torch.Tensor,
-    A_tensor_scale: float,
+    alpha: torch.Tensor,
     B_batched: torch.Tensor,
     SFB_batched: torch.Tensor,
-    B_tensor_scale: float,
     max_M: int,
     N: int,
     K: int,
@@ -1561,20 +1560,18 @@ def gemm_nvfp4_moe(
     Args:
         A_batched: Packed FP4 activations, batched (num_experts * max_M * K // 2,).
         SFA_batched: Per-expert swizzled activation scales (concatenated).
-        A_tensor_scale: Shared tensor scale for activations.
+        alpha: Device tensor (float32, 0-dim or 1-element) = act_scale * weight_scale.
         B_batched: Packed FP4 weights, batched (num_experts * N * K // 2,).
         SFB_batched: Per-expert swizzled weight scales (concatenated).
-        B_tensor_scale: Shared tensor scale for weights.
         max_M: Max tokens per expert (all experts padded to this).
         N: Output dimension per expert.
         K: Input dimension per expert.
         num_experts: Number of experts (batch dimension L).
 
     Returns:
         Output tensor (num_experts, max_M, N) in bfloat16 with tensor scales
-        applied via the CUTLASS epilogue alpha.
+        applied via the CUTLASS epilogue alpha (device-side).
     """
-    alpha = A_tensor_scale * B_tensor_scale
     return torch.ops.bitsandbytes.gemm_nvfp4_moe(
         A_batched, B_batched, SFA_batched, SFB_batched,
         alpha, max_M, N, K, num_experts,

bitsandbytes/nn/modules.py

@@ -950,10 +950,14 @@ def _forward_batched(self, x: torch.Tensor, expert_offsets: torch.Tensor) -> tor
         A_batched = torch.cat(all_act_packed)
         SFA_batched = torch.cat(all_act_scales)
 
-        # Run batched GEMM
+        # Run batched GEMM (alpha is a device tensor for graph safety)
+        alpha_dev = torch.tensor(
+            [act_tensor_scale * self.weight_tensor_scale],
+            dtype=torch.float32, device=x.device,
+        )
         D = gemm_nvfp4_moe(
-            A_batched, SFA_batched, act_tensor_scale,
-            self.weight_packed, self.weight_scales_batched, self.weight_tensor_scale,
+            A_batched, SFA_batched, alpha_dev,
+            self.weight_packed, self.weight_scales_batched,
             max_M, N, K, num_experts,
         )
 

csrc/qutlass/gemm_nvfp4_moe_sm100.cu

@@ -151,6 +151,11 @@ struct MoeGemmState {
     // Workspace
     void* workspace_dev = nullptr;
     size_t workspace_size = 0;
+
+    // Persistent GEMM object: avoids stack allocation per call, keeps
+    // params_ alive for CUDA graph replay.  init() triggers the one-time
+    // cudaFuncSetAttribute call; run() reuses the object.
+    Gemm gemm;
 };
 
 static MoeGemmState s_state;
@@ -267,13 +272,20 @@ extern "C" int cgemm_nvfp4_moe_sm100_init(
     arguments.epilogue.thread.alpha = 1.0f;
     arguments.epilogue.thread.beta = 0.0f;
 
-    Gemm gemm;
-    auto status = gemm.can_implement(arguments);
+    auto status = st.gemm.can_implement(arguments);
     if (status != cutlass::Status::kSuccess) {
         fprintf(stderr, "MoE GEMM can_implement failed: %d\n", (int)status);
         return -1;
     }
 
+    // Initialize the persistent Gemm object: triggers cudaFuncSetAttribute
+    // (one-time, not graph-safe) and fills internal params_ with dummy pointers.
+    status = st.gemm.initialize(arguments, st.workspace_dev, stream);
+    if (status != cutlass::Status::kSuccess) {
+        fprintf(stderr, "MoE GEMM initial initialize failed: %d\n", (int)status);
+        return -2;
+    }
+
     st.initialized = true;
     return 0;
 
@@ -324,13 +336,17 @@ extern "C" size_t cgemm_nvfp4_moe_sm100_workspace_size(
 // SFA_dev: activation scale factors (batched swizzled layout)
 // SFB_dev: weight scale factors (batched swizzled layout)
 // D_dev: output (num_experts, max_M, N_output) BF16, row-major per expert
+// alpha_dev: device pointer to float alpha (= act_scale * weight_scale)
+//
+// Graph-safe: only host-side param building + kernel launch.
+// cudaFuncSetAttribute was already called during _init.
 extern "C" int cgemm_nvfp4_moe_sm100_run(
     const void* A_dev,        // activations (packed FP4)
     const void* B_dev,        // weights (packed FP4)
     const void* SFA_dev,      // activation scale factors
     const void* SFB_dev,      // weight scale factors
     void* D_dev,              // output (BF16)
-    float alpha,
+    const float* alpha_dev,   // device pointer to alpha scalar
     cudaStream_t stream
 ) {
 #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
@@ -363,18 +379,23 @@ extern "C" int cgemm_nvfp4_moe_sm100_run(
          static_cast<ElementD*>(D_dev), st.stride_D},
         st.hw_info
     };
-    arguments.epilogue.thread.alpha = alpha;
+    // Device-side alpha: if alpha_dev is non-null, kernel reads from device ptr.
+    // alpha_ptr takes precedence over the scalar alpha value.
+    arguments.epilogue.thread.alpha = 1.0f;  // fallback (ignored when alpha_ptr set)
+    arguments.epilogue.thread.alpha_ptr = alpha_dev;
     arguments.epilogue.thread.beta = 0.0f;
 
-    Gemm gemm;
-
-    auto status = gemm.initialize(arguments, st.workspace_dev, stream);
+    // Rebuild params from arguments (host-side only, no CUDA API calls).
+    // cudaFuncSetAttribute was already called during _init on the persistent
+    // gemm object, so we call initialize() which is idempotent for the
+    // attribute and only updates params_.
+    auto status = st.gemm.initialize(arguments, st.workspace_dev, stream);
     if (status != cutlass::Status::kSuccess) {
         fprintf(stderr, "MoE GEMM initialize failed: %d\n", (int)status);
         return -2;
     }
 
-    status = gemm.run(stream);
+    status = st.gemm.run(stream);
     if (status != cutlass::Status::kSuccess) {
         fprintf(stderr, "MoE GEMM run failed: %d\n", (int)status);
         return -3;