feat: adapt for Qwen3.6 moe inference

src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh

@@ -21,6 +21,28 @@ __device__ __forceinline__ const T *typedInputPtr(const void *ptr) {
     return reinterpret_cast<const T *>(ptr);
 }
 
+template <size_t N>
+struct InputPointerArray {
+    const void *values[N];
+};
+
+/**
+ * @brief Stores elementwise input pointers in device workspace.
+ *
+ * The pointer array is passed by value as a kernel argument. This is required
+ * for CUDA Graph capture: a captured cudaMemcpyAsync from inputs.data() would
+ * retain a pointer to a temporary host std::vector that is destroyed before
+ * graph replay.
+ */
+template <size_t N>
+INFINIOP_CUDA_KERNEL storeInputPointers(
+    const void **output,
+    InputPointerArray<N> inputs) {
+    for (size_t i = threadIdx.x; i < N; i += blockDim.x) {
+        output[i] = inputs.values[i];
+    }
+}
+
 /**
  * @brief Computes the output index in memory, accounting for strides if non-contiguous.
  *
@@ -297,8 +319,21 @@ private:
         const int8_t *info_meta_start = info.getMetaStart();
         const int8_t *d_meta_start = reinterpret_cast<int8_t *>(workspace) + input_arr_size;
 
-        // copy the input pointer array and meta to device
-        CHECK_CUDA(cudaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, cudaMemcpyHostToDevice, stream));
+        cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
+        CHECK_CUDA(cudaStreamIsCapturing(stream, &capture_status));
+        if (capture_status == cudaStreamCaptureStatusNone) {
+            CHECK_CUDA(cudaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, cudaMemcpyHostToDevice, stream));
+        } else {
+            // A captured H2D copy from a temporary std::vector would retain an
+            // invalid host pointer for replay. Kernel arguments are stored by
+            // value in the graph node instead.
+            InputPointerArray<N> input_pointers{};
+            for (size_t i = 0; i < N; ++i) {
+                input_pointers.values[i] = h_inputs_arr[i];
+            }
+            storeInputPointers<N><<<1, N, 0, stream>>>(
+                reinterpret_cast<const void **>(workspace), input_pointers);
+        }
         CHECK_CUDA(cudaMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), cudaMemcpyHostToDevice, stream));
 
         // offset/assign the pointers

src/infiniop/ops/diff/nvidia/diff_nvidia.cuh

@@ -1,8 +1,11 @@
 #ifndef __DIFF_NVIDIA_H__
 #define __DIFF_NVIDIA_H__
 
+#include "../../../../utils.h"
 #include "../../../operator.h"
 #include <cstddef>
+#include <utility>
+#include <vector>
 
 namespace op::diff::nvidia {
 

src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu

@@ -403,6 +403,7 @@ infiniStatus_t launch_cutlass_gemm_grouped_device_meta(int problem_count,
                                                        int64_t *d_ldb,
                                                        int64_t *d_ldc,
                                                        int64_t *d_ldd,
+                                                       bool full_occupancy_wave,
                                                        cudaStream_t stream) {
     if (problem_count == 0) {
         return INFINI_STATUS_SUCCESS;
@@ -436,7 +437,15 @@ infiniStatus_t launch_cutlass_gemm_grouped_device_meta(int problem_count,
         cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>::GemmKernel;
     using Gemm = cutlass::gemm::device::GemmGrouped<GemmKernel>;
 
-    const int threadblock_count = std::min(Gemm::sufficient(), problem_count);
+    const int occupancy_wave = Gemm::sufficient();
+    // Decode routes to only top-k experts, but each expert still contains many
+    // N-dimension tiles. A full persistent wave lets CUTLASS distribute those
+    // tiles across the GPU instead of leaving most SMs idle. Prefill retains
+    // the previous problem-count cap to avoid scheduler overhead from excess
+    // threadblocks when most experts receive no tokens.
+    const int threadblock_count = full_occupancy_wave
+                                    ? occupancy_wave
+                                    : std::min(occupancy_wave, problem_count);
     if (threadblock_count <= 0) {
         return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED;
     }
@@ -547,7 +556,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
         // A D2H count copy and stream sync cannot be captured for replay.
         CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
             topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
-            grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
+            grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, stream));
 
         swiglu_kernel<T><<<(topk * intermediate_size + 255) / 256, 256, 0, stream>>>(gate_up, activated, topk, intermediate_size);
 
@@ -560,7 +569,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
             output_permutation, pairs, num_experts, hidden_size, intermediate_size, block_size);
         CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
             topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
-            grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
+            grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, stream));
 
         apply_shuffle_mul_sum_kernel<T><<<1, std::min(hidden_size, 1024), 0, stream>>>(
             expert_out, reinterpret_cast<T *>(output), output_permutation,
@@ -604,7 +613,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
     // avoids making the runtime problem count depend on device routing data.
     CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
         num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
-        grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
+        grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, stream));
 
     swiglu_kernel<T><<<(max_num_tokens_padded * intermediate_size + 255) / 256, 256, 0, stream>>>(
         gate_up, activated, max_num_tokens_padded, intermediate_size);
@@ -615,7 +624,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
         w2_t, expert_out, num_experts, hidden_size, intermediate_size);
     CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
         num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
-        grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
+        grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, stream));
 
     apply_shuffle_mul_sum_kernel<T><<<num_tokens, std::min(hidden_size, 1024), 0, stream>>>(
         expert_out, reinterpret_cast<T *>(output), output_permutation,