submission 158766

  ⋯ 9 unchanged lines
  #include <cuda_runtime.h>
  #include "cutlass/cutlass.h"
  #include "cute/tensor.hpp"
- #include "cutlass/tensor_ref.h"
  #include "cutlass/gemm/dispatch_policy.hpp"
  #include "cutlass/gemm/collective/collective_builder.hpp"
  #include "cutlass/epilogue/collective/collective_builder.hpp"
  ⋯ 6 unchanged lines
   
  #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
   
- // NVFP4 with float_ue4m3_t scale factors (VS=16, block16)
  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
  using LayoutATag = cutlass::layout::RowMajor;
- constexpr int AlignmentA = 32;
+ constexpr int AlignmentA = 64;
   
  using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
  using LayoutBTag = cutlass::layout::ColumnMajor;
- constexpr int AlignmentB = 32;
+ constexpr int AlignmentB = 64;
   
  using ElementD = cutlass::half_t;
- using ElementC = void;  // No bias matrix - pure D = A*B
+ using ElementC = void;
  using LayoutCTag = cutlass::layout::RowMajor;
  using LayoutDTag = cutlass::layout::RowMajor;
  constexpr int AlignmentD = 16;
- constexpr int AlignmentC = 1;  // void type
+ constexpr int AlignmentC = 1;
   
  using ElementAccumulator = float;
+ using ElementCompute = cutlass::half_t;  // Use half for epilogue compute
+  
  using ArchTag = cutlass::arch::Sm100;
  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
   
  using MmaTileShape = Shape<_128, _64, _256>;
  using ClusterShape = Shape<_1, _2, _1>;
   
- using KernelSchedule  = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100;
+ using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100;
  using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized1Sm;
   
  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
      ArchTag, OperatorClass,
      MmaTileShape, ClusterShape,
      cutlass::epilogue::collective::EpilogueTileAuto,
-     ElementAccumulator, ElementAccumulator,
+     ElementAccumulator, ElementCompute,
      ElementC, LayoutCTag, AlignmentC,
      ElementD, LayoutDTag, AlignmentD,
      EpilogueSchedule
  ⋯ 17 unchanged lines
      void>;
   
  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  
  using StrideA = typename Gemm::GemmKernel::StrideA;
  using StrideB = typename Gemm::GemmKernel::StrideB;
  using StrideD = typename Gemm::GemmKernel::StrideD;
  ⋯ 1 unchanged lines
  using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
  using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
   
- // Cached state for repeated calls
- static thread_local Gemm* cached_gemm = nullptr;
- static thread_local void* cached_workspace = nullptr;
- static thread_local int cached_m = 0, cached_n = 0, cached_k = 0, cached_l = 0;
- static thread_local StrideA cached_stride_A;
- static thread_local StrideB cached_stride_B;
- static thread_local StrideD cached_stride_D;
- static thread_local LayoutSFA cached_layout_SFA;
- static thread_local LayoutSFB cached_layout_SFB;
-  
  void nvfp4_gemm(
      torch::Tensor a, torch::Tensor b,
      torch::Tensor sfa_perm, torch::Tensor sfb_perm,
      torch::Tensor c, int m, int n, int k, int l)
  {
-     bool dims_changed = (m != cached_m || n != cached_n || k != cached_k || l != cached_l);
-     
-     if (dims_changed || cached_gemm == nullptr) {
-         cached_stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, l});
-         cached_stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, l});
-         cached_stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, l});
-         cached_layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, l));
-         cached_layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, l));
-     }
+     StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, l});
+     StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, l});
+     StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, l});
   
-     auto* a_ptr = reinterpret_cast<typename ElementA::DataType*>(a.data_ptr());
-     auto* b_ptr = reinterpret_cast<typename ElementB::DataType*>(b.data_ptr());
-     auto* sfa_ptr = reinterpret_cast<typename ElementA::ScaleFactorType*>(sfa_perm.data_ptr());
-     auto* sfb_ptr = reinterpret_cast<typename ElementB::ScaleFactorType*>(sfb_perm.data_ptr());
-     auto* d_ptr = reinterpret_cast<ElementD*>(c.data_ptr());
+     LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, l));
+     LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, l));
   
      typename Gemm::Arguments arguments{
          cutlass::gemm::GemmUniversalMode::kGemm,
          {m, n, k, l},
-         {a_ptr, cached_stride_A, b_ptr, cached_stride_B, sfa_ptr, cached_layout_SFA, sfb_ptr, cached_layout_SFB},
-         {{1.0f, 0.0f}, nullptr, {}, d_ptr, cached_stride_D}
+         {reinterpret_cast<typename ElementA::DataType*>(a.data_ptr()), stride_A,
+          reinterpret_cast<typename ElementB::DataType*>(b.data_ptr()), stride_B,
+          reinterpret_cast<typename ElementA::ScaleFactorType*>(sfa_perm.data_ptr()), layout_SFA,
+          reinterpret_cast<typename ElementB::ScaleFactorType*>(sfb_perm.data_ptr()), layout_SFB},
+         {{cutlass::half_t(1.0f), cutlass::half_t(0.0f)}, nullptr, {}, 
+          reinterpret_cast<ElementD*>(c.data_ptr()), stride_D}
      };
-     
-     if (dims_changed || cached_gemm == nullptr) {
-         delete cached_gemm;
-         cached_gemm = new Gemm();
-         
-         size_t workspace_size = Gemm::get_workspace_size(arguments);
-         if (workspace_size > 0) {
-             cudaMalloc(&cached_workspace, workspace_size);
-         } else {
-             cached_workspace = nullptr;
-         }
-         
-         auto status = cached_gemm->initialize(arguments, cached_workspace);
-         TORCH_CHECK(status == cutlass::Status::kSuccess, 
-             "CUTLASS init failed: ", cutlass::cutlassGetStatusString(status));
-         
-         cached_m = m; cached_n = n; cached_k = k; cached_l = l;
-     } else {
-         auto status = cached_gemm->update(arguments, cached_workspace);
-         TORCH_CHECK(status == cutlass::Status::kSuccess, 
-             "CUTLASS update failed: ", cutlass::cutlassGetStatusString(status));
-     }
   
-     auto status = cached_gemm->run();
-     TORCH_CHECK(status == cutlass::Status::kSuccess, 
-         "CUTLASS run failed: ", cutlass::cutlassGetStatusString(status));
+     Gemm gemm;
+     size_t workspace_size = Gemm::get_workspace_size(arguments);
+     auto workspace = torch::empty({static_cast<long>(workspace_size)}, 
+         torch::TensorOptions().dtype(torch::kUInt8).device(a.device()));
+  
+     gemm.initialize(arguments, workspace.data_ptr());
+     gemm.run();
  }
   
  #else
- void nvfp4_gemm(torch::Tensor a, torch::Tensor b,
-                 torch::Tensor sfa_perm, torch::Tensor sfb_perm,
-                 torch::Tensor c, int m, int n, int k, int l) {
+ void nvfp4_gemm(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+                 torch::Tensor, int, int, int, int) {
      TORCH_CHECK(false, "SM100 not supported");
  }
  #endif
  ⋯ 27 unchanged lines
          "-arch=sm_100a",
          "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
          "-O3",
+         "--use_fast_math",
      ],
      verbose=False,
  )
   
  def custom_kernel(data: input_t) -> output_t:
      a, b, sfa, sfb, sfa_perm, sfb_perm, c = data
-     m = int(c.shape[0])
-     n = int(c.shape[1])
-     l = int(c.shape[2])
-     k = int(a.shape[1] * 2)
+     m, n, l = c.shape[0], c.shape[1], c.shape[2]
+     k = a.shape[1] * 2
      module.nvfp4_gemm(a, b, sfa_perm, sfb_perm, c, m, n, k, l)
      return c
   No newline at end of file