fix: Add stream parameter to cgemm_nvfp4 for CUDA graph support

The kernel launch now uses the caller's stream via <<<grid, threads, 0,
stream>>>. The Python dispatch passes _get_tensor_stream(A_packed).
This enables CUDA graph capture for accurate benchmarking.

Author: TimDettmers

bitsandbytes/backends/cuda/ops.py

@@ -928,6 +928,7 @@ def _(
             ct.c_int(M),
             ct.c_int(N),
             ct.c_int(K),
+            _get_tensor_stream(A_packed),
         )
 
     # Apply tensor scales (the GEMM kernel operates on raw quantized values)

csrc/kernels_nvfp4_sm120.cu

@@ -485,13 +485,13 @@ __global__ void kGemmNVFP4_simple(
 // ============================================================================
 extern "C" void cgemm_nvfp4(
     const unsigned char* A, const unsigned char* B, const unsigned char* SFA, const unsigned char* SFB, float* D, int M,
-    int N, int K
+    int N, int K, cudaStream_t stream
 ) {
     int num_m_blocks = (M + BLOCK_M_DIM - 1) / BLOCK_M_DIM;
     int num_n_blocks = (N + BLOCK_N_DIM - 1) / BLOCK_N_DIM;
 
     dim3 grid(num_n_blocks, num_m_blocks);
     int threads_per_block = WARPS_PER_BLOCK * 32; // 256
 
-    kGemmNVFP4_smem<<<grid, threads_per_block>>>(A, B, SFA, SFB, D, M, N, K);
+    kGemmNVFP4_smem<<<grid, threads_per_block, 0, stream>>>(A, B, SFA, SFB, D, M, N, K);
 }