submission 189094

  import torch
- from torch.utils.cpp_extension import load_inline
-  
-  
- cpp_src = r"""
- #include <torch/extension.h>
- #include <ATen/ATen.h>
- #include <ATen/ops/_scaled_mm.h>
-  
- void silu_mul_cuda(const at::Half* g1, const at::Half* g2, at::Half* out, int64_t count, int64_t out_stride);
-  
- torch::Tensor to_blocked(torch::Tensor input) {
-     auto rows = input.size(0);
-     auto cols = input.size(1);
-     auto n_row_blocks = (rows + 127) / 128;
-     auto n_col_blocks = (cols + 3) / 4;
-     auto padded = input.contiguous();
-     auto blocks = padded.view({n_row_blocks, 128, n_col_blocks, 4}).permute({0, 2, 1, 3});
-     auto rearranged = blocks.reshape({-1, 4, 32, 4}).transpose(1, 2).reshape({-1, 32, 16});
-     return rearranged.reshape({-1}).contiguous();
- }
-  
- torch::Tensor fused_dual_gemm(torch::Tensor a, torch::Tensor b1, torch::Tensor b2, torch::Tensor sfa, torch::Tensor sfb1, torch::Tensor sfb2, torch::Tensor c) {
-     auto out = c.contiguous();
-     auto l = out.size(2);
-     auto out_stride = out.size(2);
-     for (int64_t l_idx = 0; l_idx < l; ++l_idx) {
-         auto a_l = a.select(2, l_idx);
-         auto b1_l = b1.select(2, l_idx);
-         auto b2_l = b2.select(2, l_idx);
-         auto scale_a = to_blocked(sfa.select(2, l_idx));
-         auto scale_b1 = to_blocked(sfb1.select(2, l_idx));
-         auto scale_b2 = to_blocked(sfb2.select(2, l_idx));
-         auto g1 = at::_scaled_mm(
-             a_l,
-             b1_l.transpose(0, 1),
-             scale_a,
-             scale_b1,
-             c10::optional<at::Tensor>(),
-             c10::optional<at::Tensor>(),
-             c10::optional<at::ScalarType>(at::kHalf),
-             false
-         );
-         auto g2 = at::_scaled_mm(
-             a_l,
-             b2_l.transpose(0, 1),
-             scale_a,
-             scale_b2,
-             c10::optional<at::Tensor>(),
-             c10::optional<at::Tensor>(),
-             c10::optional<at::ScalarType>(at::kHalf),
-             false
-         );
-         auto g1_c = g1.contiguous();
-         auto g2_c = g2.contiguous();
-         int64_t count = g1_c.numel();
-         auto out_ptr = reinterpret_cast<at::Half*>(out.data_ptr()) + l_idx;
-         silu_mul_cuda(
-             reinterpret_cast<const at::Half*>(g1_c.data_ptr()),
-             reinterpret_cast<const at::Half*>(g2_c.data_ptr()),
-             out_ptr,
-             count,
-             out_stride
-         );
-     }
-     return out;
- }
-  
- PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fused_dual_gemm", &fused_dual_gemm, "nvfp4 dual gemm"); }
- """
-  
-  
- cuda_src = r"""
- #include <ATen/ATen.h>
- #include <cuda.h>
- #include <cuda_fp16.h>
- #include <cuda_runtime.h>
- #include <math.h>
- #include <stdint.h>
-  
- __global__ void silu_mul_kernel(const half* g1, const half* g2, half* out, int64_t count, int64_t out_stride) {
-     int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-     if (idx >= count) {
-         return;
-     }
-     float x = __half2float(g1[idx]);
-     float y = __half2float(g2[idx]);
-     float silu = x / (1.0f + expf(-x));
-     out[idx * out_stride] = __float2half(silu * y);
- }
-  
- void silu_mul_cuda(const at::Half* g1, const at::Half* g2, at::Half* out, int64_t count, int64_t out_stride) {
-     int threads = 256;
-     int blocks = static_cast<int>((count + threads - 1) / threads);
-     auto g1_ptr = reinterpret_cast<const half*>(g1);
-     auto g2_ptr = reinterpret_cast<const half*>(g2);
-     auto out_ptr = reinterpret_cast<half*>(out);
-     silu_mul_kernel<<<blocks, threads>>>(g1_ptr, g2_ptr, out_ptr, count, out_stride);
- }
- """
-  
-  
- ext = load_inline(
-     name="nvfp4_dual_gemm_ext",
-     cpp_sources=cpp_src,
-     cuda_sources=cuda_src,
-     functions=None,
-     with_cuda=True,
-     extra_cflags=[
-         "-O3",
-         "-std=c++17",
-     ],
-     extra_cuda_cflags=[
-         "-O3",
-         "--use_fast_math",
-         "-lineinfo",
-     ],
-     verbose=False,
- )
-  
-  
+ from typing import Tuple, Type, Optional, Union
+ import cutlass
+ import cutlass.cute as cute
+ from cutlass.cute.runtime import make_ptr
+ from cutlass.cute.nvgpu import cpasync, tcgen05
+ import cutlass.utils as utils
+ import cutlass.pipeline as pipeline
+ import cutlass.utils.blackwell_helpers as sm100_utils
+ import cutlass.utils.blockscaled_layout as blockscaled_utils
+ mma_inst_shape_k = 64
+ ab_dtype = cutlass.Float4E2M1FN
+ sf_dtype = cutlass.Float8E4M3FN
+ c_dtype = cutlass.Float16
+ sf_vec_size = 16
+ class Sm100BlockScaledDenseGemmKernel:
+     def __init__(
+         self,
+         mma_tiler_mn: Tuple[int, int],
+         cluster_shape_mn: Tuple[int, int],
+     ):
+         self.ab_dtype = cutlass.Float4E2M1FN
+         self.sf_dtype = cutlass.Float8E4M3FN
+         self.acc_dtype = cutlass.Float32
+         self.c_dtype = cutlass.Float16
+         self.sf_vec_size = 16
+         self.epilog_warp_id = (0, 1, 2, 3)
+         self.mma_warp_id = 4
+         self.tma_warp_id = 5
+         self.threads_per_cta = 32 * len((self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id))
+         self.mma_tiler_mn = mma_tiler_mn
+         self.cluster_shape_mn = cluster_shape_mn
+         self.occupancy = 1
+         self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
+         self.num_tmem_alloc_cols = 512
+         self.epilog_sync_barrier = pipeline.NamedBarrier(
+             barrier_id=1,
+             num_threads=32 * len(self.epilog_warp_id),
+         )
+         self.tmem_alloc_barrier = pipeline.NamedBarrier(
+             barrier_id=2,
+             num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
+         )
+     def _setup_attributes(self):
+         tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+             self.ab_dtype,
+             self.a_major_mode,
+             self.b_major_mode,
+             self.sf_dtype,
+             self.sf_vec_size,
+             tcgen05.CtaGroup.ONE,
+             self.mma_tiler_mn,
+         )
+         mma_inst_tile_k = 4
+         mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+         self.mma_tiler = (
+             self.mma_tiler_mn[0],
+             self.mma_tiler_mn[1],
+             mma_inst_shape_k * mma_inst_tile_k,
+         )
+         self.cta_tile_shape_mnk = (
+             self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+             self.mma_tiler[1],
+             self.mma_tiler[2],
+         )
+         self.cluster_layout_vmnk = cute.tiled_divide(
+             cute.make_layout((*self.cluster_shape_mn, 1)),
+             (tiled_mma.thr_id.shape,),
+         )
+         self.mma_inst_shape_mn_sfb = (
+             self.mma_tiler_mn[0],
+             cute.round_up(self.mma_tiler_mn[1], 128),
+         )
+         tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
+             self.ab_dtype,
+             self.a_major_mode,
+             self.b_major_mode,
+             self.sf_dtype,
+             self.sf_vec_size,
+             tcgen05.CtaGroup.ONE,
+             self.mma_inst_shape_mn_sfb,
+         )
+         self.mma_tiler_sfb = (
+             self.mma_inst_shape_mn_sfb[0],
+             self.mma_inst_shape_mn_sfb[1],
+             mma_inst_shape_k * mma_inst_tile_k,
+         )
+         self.cluster_layout_sfb_vmnk = cute.tiled_divide(
+             cute.make_layout((*self.cluster_shape_mn, 1)),
+             (tiled_mma_sfb.thr_id.shape,),
+         )
+         self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+         self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
+         self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
+         self.is_a_mcast = self.num_mcast_ctas_a > 1
+         self.is_b_mcast = self.num_mcast_ctas_b > 1
+         self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
+         self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
+             self.cta_tile_shape_mnk,
+             False,
+             self.c_layout,
+             self.c_dtype,
+         )
+         self.epi_tile_n = cute.size(self.epi_tile[1])
+         self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages(
+             tiled_mma,
+             self.mma_tiler,
+             self.a_dtype,
+             self.b_dtype,
+             self.epi_tile,
+             self.c_dtype,
+             self.c_layout,
+             self.sf_dtype,
+             self.sf_vec_size,
+             self.smem_capacity,
+             self.occupancy,
+         )
+         self.prefetch_stage = self.num_ab_stage
+         self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+             tiled_mma,
+             self.mma_tiler,
+             self.ab_dtype,
+             self.num_ab_stage,
+         )
+         self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+             tiled_mma,
+             self.mma_tiler,
+             self.ab_dtype,
+             self.num_ab_stage,
+         )
+         self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+             tiled_mma,
+             self.mma_tiler,
+             self.sf_vec_size,
+             self.num_ab_stage,
+         )
+         self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
+             tiled_mma,
+             self.mma_tiler,
+             self.sf_vec_size,
+             self.num_ab_stage,
+         )
+         self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
+             self.c_dtype,
+             self.c_layout,
+             self.epi_tile,
+             self.num_c_stage,
+         )
+     @cute.jit
+     def __call__(
+         self,
+         a_ptr: cute.Pointer,
+         b_ptr: cute.Pointer,
+         sfa_ptr: cute.Pointer,
+         sfb_ptr: cute.Pointer,
+         c_ptr: cute.Pointer,
+         m: cutlass.Int32,
+         n: cutlass.Int32,
+         k: cutlass.Int32,
+         l: cutlass.Int32,
+     ):
+         self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type
+         self.b_dtype: Type[cutlass.Numeric] = b_ptr.value_type
+         self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type
+         self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_type
+         self.a_major_mode, self.b_major_mode, self.c_layout = (
+             tcgen05.OperandMajorMode.K,
+             tcgen05.OperandMajorMode.K,
+             utils.LayoutEnum.ROW_MAJOR,
+         )
+         self._setup_attributes()
+         a_tensor = cute.make_tensor(
+             a_ptr,
+             cute.make_ordered_layout(
+                 (cute.assume(m, 32), k, l), order=(1, 0, 2)
+             ),
+         )
+         b_tensor = cute.make_tensor(
+             b_ptr,
+             cute.make_ordered_layout(
+                 (cute.assume(n, 32), k, l), order=(1, 0, 2)
+             ),
+         )
+         c_tensor = cute.make_tensor(
+             c_ptr,
+             cute.make_ordered_layout(
+                 (m, cute.assume(n, 32), l), order=(1, 0, 2)
+             )
+         )
+         sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
+             a_tensor.shape, self.sf_vec_size
+         )
+         sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
+         sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
+             b_tensor.shape, self.sf_vec_size
+         )
+         sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
+         tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+             self.ab_dtype,
+             self.a_major_mode,
+             self.b_major_mode,
+             self.sf_dtype,
+             self.sf_vec_size,
+             tcgen05.CtaGroup.ONE,
+             self.mma_tiler_mn,
+         )
+         tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
+             self.ab_dtype,
+             self.a_major_mode,
+             self.b_major_mode,
+             self.sf_dtype,
+             self.sf_vec_size,
+             tcgen05.CtaGroup.ONE,
+             self.mma_inst_shape_mn_sfb,
+         )
+         a_op = sm100_utils.cluster_shape_to_tma_atom_A(
+             self.cluster_shape_mn, tiled_mma.thr_id
+         )
+         a_smem_layout = cute.slice_(
+             self.a_smem_layout_staged, (None, None, None, 0)
+         )
+         tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
+             a_op,
+             a_tensor,
+             a_smem_layout,
+             self.mma_tiler,
+             tiled_mma,
+             self.cluster_layout_vmnk.shape,
+         )
+         b_op = sm100_utils.cluster_shape_to_tma_atom_B(
+             self.cluster_shape_mn, tiled_mma.thr_id
+         )
+         b_smem_layout = cute.slice_(
+             self.b_smem_layout_staged, (None, None, None, 0)
+         )
+         tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+             b_op,
+             b_tensor,
+             b_smem_layout,
+             self.mma_tiler,
+             tiled_mma,
+             self.cluster_layout_vmnk.shape,
+         )
+         sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
+             self.cluster_shape_mn, tiled_mma.thr_id
+         )
+         sfa_smem_layout = cute.slice_(
+             self.sfa_smem_layout_staged, (None, None, None, 0)
+         )
+         tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
+             sfa_op,
+             sfa_tensor,
+             sfa_smem_layout,
+             self.mma_tiler,
+             tiled_mma,
+             self.cluster_layout_vmnk.shape,
+             internal_type=cutlass.Int16,
+         )
+         sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
+             self.cluster_shape_mn, tiled_mma.thr_id
+         )
+         sfb_smem_layout = cute.slice_(
+             self.sfb_smem_layout_staged, (None, None, None, 0)
+         )
+         tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
+             sfb_op,
+             sfb_tensor,
+             sfb_smem_layout,
+             self.mma_tiler_sfb,           
+             tiled_mma_sfb,                
+             self.cluster_layout_sfb_vmnk.shape, 
+             internal_type=cutlass.Int16,
+         )
+         atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+         a_copy_size = cute.size_in_bytes(self.ab_dtype, a_smem_layout)
+         b_copy_size = cute.size_in_bytes(self.ab_dtype, b_smem_layout)
+         sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
+         sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
+         self.num_tma_load_bytes = (
+             a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
+         ) * atom_thr_size
+         epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
+         tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
+             cpasync.CopyBulkTensorTileS2GOp(),
+             c_tensor,
+             epi_smem_layout,
+             self.epi_tile,
+         )
+         grid = self._compute_grid(
+             c_tensor, self.cta_tile_shape_mnk, self.cluster_shape_mn
+         )
+         self.buffer_align_bytes = 1024
+         @cute.struct
+         class SharedStorage:
+             ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
+             ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
+             acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
+             tmem_dealloc_mbar_ptr: cutlass.Int64
+             tmem_holding_buf: cutlass.Int32
+             sC: cute.struct.Align[
+                 cute.struct.MemRange[
+                     self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)
+                 ],
+                 self.buffer_align_bytes,
+             ]
+             sA: cute.struct.Align[
+                 cute.struct.MemRange[
+                     self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
+                 ],
+                 self.buffer_align_bytes,
+             ]
+             sB: cute.struct.Align[
+                 cute.struct.MemRange[
+                     self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
+                 ],
+                 self.buffer_align_bytes,
+             ]
+             sSFA: cute.struct.Align[
+                 cute.struct.MemRange[
+                     self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
+                 ],
+                 self.buffer_align_bytes,
+             ]
+             sSFB: cute.struct.Align[
+                 cute.struct.MemRange[
+                     self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
+                 ],
+                 self.buffer_align_bytes,
+             ]
+         self.shared_storage = SharedStorage
+         self.kernel(
+             tiled_mma,
+             tiled_mma_sfb,                  
+             tma_atom_a,
+             tma_tensor_a,
+             tma_atom_b,
+             tma_tensor_b,
+             tma_atom_sfa,
+             tma_tensor_sfa,
+             tma_atom_sfb,
+             tma_tensor_sfb,
+             tma_atom_c,
+             tma_tensor_c,
+             self.cluster_layout_vmnk,
+             self.cluster_layout_sfb_vmnk,   
+             self.a_smem_layout_staged,
+             self.b_smem_layout_staged,
+             self.sfa_smem_layout_staged,
+             self.sfb_smem_layout_staged,
+             self.c_smem_layout_staged,
+             self.epi_tile,
+         ).launch(
+             grid=grid,
+             block=[self.threads_per_cta, 1, 1],
+             cluster=(*self.cluster_shape_mn, 1),
+             smem=self.shared_storage.size_in_bytes(),
+         )
+         return
+     @cute.kernel
+     def kernel(
+         self,
+         tiled_mma: cute.TiledMma,
+         tiled_mma_sfb: cute.TiledMma,       
+         tma_atom_a: cute.CopyAtom,
+         mA_mkl: cute.Tensor,
+         tma_atom_b: cute.CopyAtom,
+         mB_nkl: cute.Tensor,
+         tma_atom_sfa: cute.CopyAtom,
+         mSFA_mkl: cute.Tensor,
+         tma_atom_sfb: cute.CopyAtom,
+         mSFB_nkl: cute.Tensor,
+         tma_atom_c: Optional[cute.CopyAtom],
+         mC_mnl: cute.Tensor,
+         cluster_layout_vmnk: cute.Layout,
+         cluster_layout_sfb_vmnk: cute.Layout, 
+         a_smem_layout_staged: cute.ComposedLayout,
+         b_smem_layout_staged: cute.ComposedLayout,
+         sfa_smem_layout_staged: cute.Layout,
+         sfb_smem_layout_staged: cute.Layout,
+         c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None],
+         epi_tile: cute.Tile,
+     ):
+         warp_idx = cute.arch.warp_idx()
+         warp_idx = cute.arch.make_warp_uniform(warp_idx)
+         if warp_idx == self.tma_warp_id:
+             cpasync.prefetch_descriptor(tma_atom_a)
+             cpasync.prefetch_descriptor(tma_atom_b)
+             cpasync.prefetch_descriptor(tma_atom_sfa)
+             cpasync.prefetch_descriptor(tma_atom_sfb)
+             cpasync.prefetch_descriptor(tma_atom_c)
+         bidx, bidy, bidz = cute.arch.block_idx()
+         mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+         cta_rank_in_cluster = cute.arch.make_warp_uniform(
+             cute.arch.block_idx_in_cluster()
+         )
+         block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
+             cta_rank_in_cluster
+         )
+         block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
+             cta_rank_in_cluster
+         )
+         cta_coord = (bidx, bidy, bidz)
+         mma_tile_coord_mnl = (
+             cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
+             cta_coord[1],
+             cta_coord[2],
+         )
+         tidx, _, _ = cute.arch.thread_idx()
+         smem = utils.SmemAllocator()
+         storage = smem.allocate(self.shared_storage)
+         ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+         num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+         ab_pipeline_consumer_group = pipeline.CooperativeGroup(
+             pipeline.Agent.Thread, num_tma_producer
+         )
+         ab_pipeline = pipeline.PipelineTmaAsync.create(
+             barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
+             num_stages=self.num_ab_stage,
+             producer_group=ab_pipeline_producer_group,
+             consumer_group=ab_pipeline_consumer_group,
+             tx_count=self.num_tma_load_bytes,
+             cta_layout_vmnk=cluster_layout_vmnk,
+             defer_sync=True,
+         )
+         acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+         num_acc_consumer_threads = len(self.epilog_warp_id)
+         acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+             pipeline.Agent.Thread, num_acc_consumer_threads,
+         )
+         acc_pipeline = pipeline.PipelineUmmaAsync.create(
+             barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
+             num_stages=self.num_acc_stage,
+             producer_group=acc_pipeline_producer_group,
+             consumer_group=acc_pipeline_consumer_group,
+             cta_layout_vmnk=cluster_layout_vmnk,
+         )
+         tmem = utils.TmemAllocator(
+             storage.tmem_holding_buf,
+             barrier_for_retrieve=self.tmem_alloc_barrier,
+             allocator_warp_id=self.epilog_warp_id[0],
+         )
+         cute.arch.cluster_arrive_relaxed()
+         sC = storage.sC.get_tensor(
+             c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
+         )
+         sA = storage.sA.get_tensor(
+             a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
+         )
+         sB = storage.sB.get_tensor(
+             b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
+         )
+         sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+         sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
+         a_full_mcast_mask = cpasync.create_tma_multicast_mask(
+             cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+         )
+         b_full_mcast_mask = cpasync.create_tma_multicast_mask(
+             cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+         )
+         sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
+             cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+         )
+         sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
+             cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
+         )
+         gA_mkl = cute.local_tile(
+             mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+         )
+         gB_nkl = cute.local_tile(
+             mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+         )
+         gSFA_mkl = cute.local_tile(
+             mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+         )
+         gSFB_nkl = cute.local_tile(
+             mSFB_nkl,
+             cute.slice_(self.mma_tiler_sfb, (0, None, None)),
+             (None, None, None),
+         )
+         gC_mnl = cute.local_tile(
+             mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
+         )
+         k_block_cnt = cute.size(gA_mkl, mode=[3])
+         thr_mma = tiled_mma.get_slice(0)
+         thr_mma_sfb = tiled_mma_sfb.get_slice(0) 
+         thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+         thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
+         tCgA = thr_mma.partition_A(gA_mkl)
+         tCgB = thr_mma.partition_B(gB_nkl)
+         tCgSFA = thr_mma.partition_A(gSFA_mkl)
+         tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
+         tCgC = thr_mma.partition_C(gC_mnl)
+         a_cta_layout = cute.make_layout(
+             cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
+         )
+         tAsA, tAgA = cpasync.tma_partition(
+             tma_atom_a,
+             block_in_cluster_coord_vmnk[2],
+             a_cta_layout,
+             cute.group_modes(sA, 0, 3),
+             cute.group_modes(tCgA, 0, 3),
+         )
+         b_cta_layout = cute.make_layout(
+             cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
+         )
+         tBsB, tBgB = cpasync.tma_partition(
+             tma_atom_b,
+             block_in_cluster_coord_vmnk[1],
+             b_cta_layout,
+             cute.group_modes(sB, 0, 3),
+             cute.group_modes(tCgB, 0, 3),
+         )
+         sfa_cta_layout = a_cta_layout
+         tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
+             tma_atom_sfa,
+             block_in_cluster_coord_vmnk[2],
+             sfa_cta_layout,
+             cute.group_modes(sSFA, 0, 3),
+             cute.group_modes(tCgSFA, 0, 3),
+         )
+         tAsSFA = cute.filter_zeros(tAsSFA)
+         tAgSFA = cute.filter_zeros(tAgSFA)
+         sfb_cta_layout = cute.make_layout(
+             cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
+         )
+         tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition(
+             tma_atom_sfb,
+             block_in_cluster_coord_sfb_vmnk[1],
+             sfb_cta_layout,
+             cute.group_modes(sSFB, 0, 3),
+             cute.group_modes(tCgSFB, 0, 3),
+         )
+         tBsSFB = cute.filter_zeros(tBsSFB)
+         tBgSFB = cute.filter_zeros(tBgSFB)
+         tCrA = tiled_mma.make_fragment_A(sA)
+         tCrB = tiled_mma.make_fragment_B(sB)
+         acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+         tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+         cute.arch.cluster_wait()
+         if warp_idx == self.tma_warp_id:
+             ab_producer_state = pipeline.make_pipeline_state(
+                 pipeline.PipelineUserType.Producer, self.num_ab_stage
+             )
+             tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+             tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+             tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+             slice_n = mma_tile_coord_mnl[1]
+             if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
+                 slice_n = mma_tile_coord_mnl[1] // 2
+             tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])]
+             for prefetch_tile in cutlass.range(0, self.prefetch_stage, unroll=1):
+                 cute.prefetch(
+                     tma_atom_a,
+                     tAgA_slice[(None, prefetch_tile)]
+                 )
+                 cute.prefetch(
+                     tma_atom_b,
+                     tBgB_slice[(None, prefetch_tile)]
+                 )
+                 cute.prefetch(
+                     tma_atom_sfa,
+                     tAgSFA_slice[(None, prefetch_tile)]
+                 )
+                 cute.prefetch(
+                     tma_atom_sfb,
+                     tBgSFB_slice[(None, prefetch_tile)]
+                 )
+             peek_ab_empty_status = ab_pipeline.producer_try_acquire(
+                 ab_producer_state
+             )
+             for k_block_idx in cutlass.range(0, k_block_cnt, 1, unroll=1):
+                 ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
+                 cute.copy(
+                     tma_atom_a,
+                     tAgA_slice[(None, ab_producer_state.count)],
+                     tAsA[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=a_full_mcast_mask,
+                 )
+                 cute.copy(
+                     tma_atom_b,
+                     tBgB_slice[(None, ab_producer_state.count)],
+                     tBsB[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=b_full_mcast_mask,
+                 )
+                 cute.copy(
+                     tma_atom_sfa,
+                     tAgSFA_slice[(None, ab_producer_state.count)],
+                     tAsSFA[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=sfa_full_mcast_mask,
+                 )
+                 cute.copy(
+                     tma_atom_sfb,
+                     tBgSFB_slice[(None, ab_producer_state.count)],
+                     tBsSFB[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=sfb_full_mcast_mask,
+                 )
+                 if k_block_idx < k_block_cnt - self.prefetch_stage:
+                     next_k_idx = ab_producer_state.count + self.prefetch_stage
+                     cute.prefetch(
+                         tma_atom_a,
+                         tAgA_slice[(None, next_k_idx)]
+                     )
+                     cute.prefetch(
+                         tma_atom_b,
+                         tBgB_slice[(None, next_k_idx)]
+                     )
+                     cute.prefetch(
+                         tma_atom_sfa,
+                         tAgSFA_slice[(None, next_k_idx)]
+                     )
+                     cute.prefetch(
+                         tma_atom_sfb,
+                         tBgSFB_slice[(None, next_k_idx)]
+                     )
+                 ab_producer_state.advance()
+                 if ab_producer_state.count < k_block_cnt:
+                     peek_ab_empty_status = ab_pipeline.producer_try_acquire(
+                         ab_producer_state
+                     )
+             ab_pipeline.producer_tail(ab_producer_state)
+         elif warp_idx == self.mma_warp_id:
+             tmem.wait_for_alloc()
+             acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) 
+             tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+             sfa_tmem_ptr = cute.recast_ptr(
+                 acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
+                 dtype=self.sf_dtype,
+             )
+             tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+                 tiled_mma,
+                 self.mma_tiler,
+                 self.sf_vec_size,
+                 cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+             )
+             tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) 
+             sfb_tmem_ptr = cute.recast_ptr(
+                 acc_tmem_ptr
+                 + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+                 + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
+                 dtype=self.sf_dtype,
+             )
+             tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+                 tiled_mma,
+                 self.mma_tiler,
+                 self.sf_vec_size,
+                 cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
+             )
+             tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) 
+             tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = (
+                 self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
+             )
+             tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = (
+                 self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
+             )
+             ab_consumer_state = pipeline.make_pipeline_state(
+                 pipeline.PipelineUserType.Consumer, self.num_ab_stage
+             )
+             acc_producer_state = pipeline.make_pipeline_state(
+                 pipeline.PipelineUserType.Producer, self.num_acc_stage
+             )
+             peek_ab_full_status = ab_pipeline.consumer_try_wait(
+                 ab_consumer_state
+             )
+             tCtSFB_mma = tCtSFB
+             if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
+                 offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
+                 shifted_ptr = cute.recast_ptr(
+                     acc_tmem_ptr
+                     + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+                     + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+                     + offset,
+                     dtype=self.sf_dtype,
+                 )
+                 tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
+             tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+             for _ in range(k_block_cnt):
+                 ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
+                 s2t_stage_coord = (
+                     None,
+                     None,
+                     None,
+                     None,
+                     ab_consumer_state.index,
+                 )
+                 tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
+                 tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
+                 cute.copy(
+                     tiled_copy_s2t_sfa,
+                     tCsSFA_compact_s2t_staged,
+                     tCtSFA_compact_s2t,
+                 )
+                 cute.copy(
+                     tiled_copy_s2t_sfb,
+                     tCsSFB_compact_s2t_staged,
+                     tCtSFB_compact_s2t,
+                 )
+                 num_kphases = cute.size(tCrA, mode=[2])
+                 for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                     kphase_coord = (
+                         None,
+                         None,
+                         kphase_idx,
+                         ab_consumer_state.index,
+                     )
+                     sf_kphase_coord = (None, None, kphase_idx)
+                     tiled_mma.set(
+                         tcgen05.Field.SFA,
+                         tCtSFA[sf_kphase_coord].iterator,
+                     )
+                     tiled_mma.set(
+                         tcgen05.Field.SFB,
+                         tCtSFB_mma[sf_kphase_coord].iterator,
+                     )
+                     cute.gemm(
+                         tiled_mma,
+                         tCtAcc,
+                         tCrA[kphase_coord],
+                         tCrB[kphase_coord],
+                         tCtAcc,
+                     )
+                     tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                 ab_pipeline.consumer_release(ab_consumer_state)
+                 ab_consumer_state.advance()
+                 if ab_consumer_state.count < k_block_cnt:
+                     peek_ab_full_status = ab_pipeline.consumer_try_wait(
+                         ab_consumer_state
+                     )
+             acc_pipeline.producer_commit(acc_producer_state)
+         elif warp_idx in self.epilog_warp_id:
+             tmem.allocate(self.num_tmem_alloc_cols)
+             tmem.wait_for_alloc()
+             acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+             tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+             tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition(
+                 tidx, tCtAcc, tCgC, epi_tile
+             )
+             tTR_rC = cute.make_fragment(tTR_rAcc.shape, self.c_dtype)
+             tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
+                 tiled_copy_t2r, tTR_rC, tidx, sC
+             )
+             tma_atom_c, bSG_sC, bSG_gC = self.epilog_gmem_copy_and_partition(
+                 tidx, tma_atom_c, tCgC, epi_tile, sC
+             )
+             bSG_gC = bSG_gC[(None, None, None, *mma_tile_coord_mnl)]
+             acc_consumer_state = pipeline.make_pipeline_state(
+                 pipeline.PipelineUserType.Consumer, self.num_acc_stage
+             )
+             acc_pipeline.consumer_wait(acc_consumer_state)
+             tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+             bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
+             subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+             for subtile_idx in range(subtile_cnt):
+                 tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
+                 cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
+                 tRS_rC.store(tTR_rAcc.load().to(self.c_dtype))
+                 cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, subtile_idx)])
+                 cute.arch.fence_view_async_shared()
+                 if warp_idx == self.epilog_warp_id[0]:
+                     cute.copy(
+                         tma_atom_c,
+                         bSG_sC[(None, subtile_idx)],
+                         bSG_gC[(None, subtile_idx)],
+                     )
+             tmem.relinquish_alloc_permit()
+             tmem.free(acc_tmem_ptr)
+     def mainloop_s2t_copy_and_partition(
+         self,
+         sSF: cute.Tensor,
+         tSF: cute.Tensor,
+     ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+         tCsSF_compact = cute.filter_zeros(sSF)
+         tCtSF_compact = cute.filter_zeros(tSF)
+         copy_atom_s2t = cute.make_copy_atom(
+             tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
+             self.sf_dtype,
+         )
+         tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
+         thr_copy_s2t = tiled_copy_s2t.get_slice(0)
+         tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
+         tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+             tiled_copy_s2t, tCsSF_compact_s2t_
+         )
+         tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
+         return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
+     def epilog_tmem_copy_and_partition(
+         self,
+         tidx: cutlass.Int32,
+         tAcc: cute.Tensor,
+         gC_mnl: cute.Tensor,
+         epi_tile: cute.Tile,
+     ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+         copy_atom_t2r = sm100_utils.get_tmem_load_op(
+             self.cta_tile_shape_mnk,
+             self.c_layout,
+             self.c_dtype,
+             self.acc_dtype,
+             epi_tile,
+             False
+         )
+         tAcc_epi = cute.flat_divide(
+             tAcc[((None, None), 0, 0)],
+             epi_tile,
+         )
+         tiled_copy_t2r = tcgen05.make_tmem_copy(
+             copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]
+         )
+         thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+         tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
+         gC_mnl_epi = cute.flat_divide(
+             gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
+         )
+         tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
+         tTR_rAcc = cute.make_fragment(
+             tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
+         )
+         return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
+     def epilog_smem_copy_and_partition(
+         self,
+         tiled_copy_t2r: cute.TiledCopy,
+         tTR_rC: cute.Tensor,
+         tidx: cutlass.Int32,
+         sC: cute.Tensor,
+     ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+         copy_atom_r2s = sm100_utils.get_smem_store_op(
+             self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r
+         )
+         tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+         thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+         tRS_sC = thr_copy_r2s.partition_D(sC)
+         tRS_rC = tiled_copy_r2s.retile(tTR_rC)
+         return tiled_copy_r2s, tRS_rC, tRS_sC
+     def epilog_gmem_copy_and_partition(
+         self,
+         tidx: cutlass.Int32,
+         atom: Union[cute.CopyAtom, cute.TiledCopy],
+         gC_mnl: cute.Tensor,
+         epi_tile: cute.Tile,
+         sC: cute.Tensor,
+     ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
+         gC_epi = cute.flat_divide(
+             gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
+         )
+         tma_atom_c = atom
+         sC_for_tma_partition = cute.group_modes(sC, 0, 2)
+         gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
+         bSG_sC, bSG_gC = cpasync.tma_partition(
+             tma_atom_c,
+             0,
+             cute.make_layout(1),
+             sC_for_tma_partition,
+             gC_for_tma_partition,
+         )
+         return tma_atom_c, bSG_sC, bSG_gC
+     @staticmethod
+     def _compute_stages(
+         tiled_mma: cute.TiledMma,
+         mma_tiler_mnk: Tuple[int, int, int],
+         a_dtype: Type[cutlass.Numeric],
+         b_dtype: Type[cutlass.Numeric],
+         epi_tile: cute.Tile,
+         c_dtype: Type[cutlass.Numeric],
+         c_layout: utils.LayoutEnum,
+         sf_dtype: Type[cutlass.Numeric],
+         sf_vec_size: int,
+         smem_capacity: int,
+         occupancy: int,
+     ) -> Tuple[int, int, int]:
+         num_acc_stage = 1
+         num_c_stage = 2
+         a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
+             tiled_mma,
+             mma_tiler_mnk,
+             a_dtype,
+             1,  
+         )
+         b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
+             tiled_mma,
+             mma_tiler_mnk,
+             b_dtype,
+             1,  
+         )
+         sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
+             tiled_mma,
+             mma_tiler_mnk,
+             sf_vec_size,
+             1,  
+         )
+         sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
+             tiled_mma,
+             mma_tiler_mnk,
+             sf_vec_size,
+             1,  
+         )
+         c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(
+             c_dtype,
+             c_layout,
+             epi_tile,
+             1,
+         )
+         ab_bytes_per_stage = (
+             cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
+             + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
+             + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+             + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one)
+         )
+         mbar_helpers_bytes = 1024
+         c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
+         c_bytes = c_bytes_per_stage * num_c_stage
+         num_ab_stage = (
+             smem_capacity - (mbar_helpers_bytes + c_bytes)
+         ) // ab_bytes_per_stage
+         num_c_stage += (
+             smem_capacity
+             -  ab_bytes_per_stage * num_ab_stage
+             -  (mbar_helpers_bytes + c_bytes)
+         ) // (c_bytes_per_stage)
+         return num_acc_stage, num_ab_stage, num_c_stage
+     @staticmethod
+     def _compute_grid(
+         c: cute.Tensor,
+         cta_tile_shape_mnk: Tuple[int, int, int],
+         cluster_shape_mn: Tuple[int, int],
+     ) -> Tuple[int, int, int]:
+         grid = (
+                 cute.ceil_div(c.layout.shape[0], cta_tile_shape_mnk[0]),
+                 cute.ceil_div(c.layout.shape[1], cta_tile_shape_mnk[1]),
+                 c.layout.shape[2],
+         )
+         return grid
+ _compiled_kernel_cache = {}
+ def compile_kernel(shape_key):
+     cached = _compiled_kernel_cache.get(shape_key)
+     if cached is not None:
+         return cached
+     a_ptr = make_ptr(
+         ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
+     )
+     b_ptr = make_ptr(
+         ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
+     )
+     c_ptr = make_ptr(
+         c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
+     )
+     sfa_ptr = make_ptr(
+         sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
+     )
+     sfb_ptr = make_ptr(
+         sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
+     )
+     my_kernel = Sm100BlockScaledDenseGemmKernel((128, 64), (1, 1))
+     compiled = cute.compile(my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, 0, 0, 0, 0)
+     _compiled_kernel_cache[shape_key] = compiled
+     return compiled
  def custom_kernel(data):
-     a, b1, b2, sfa, sfb1, sfb2, sfa_p, sfb1_p, sfb2_p, c = data
-     return ext.fused_dual_gemm(a, b1, b2, sfa, sfb1, sfb2, c)
-  
-  
+     a, b1, b2, _, _, _, sfa_permuted, sfb1_permuted, sfb2_permuted, c = data
+     c = c.contiguous()
+     m, k, l = a.shape
+     n, _, _ = b1.shape
+     # float4_e2m1fn_x2 为打包表示，逻辑K需要乘2
+     k = k * 2
+     shape_key = "default"
+     compiled_func = compile_kernel(shape_key)
+     g1 = c
+     g2 = torch.empty_like(c)
+     a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     b1_ptr = make_ptr(ab_dtype, b1.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     b2_ptr = make_ptr(ab_dtype, b2.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     g1_ptr = make_ptr(c_dtype, g1.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     g2_ptr = make_ptr(c_dtype, g2.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     sfa_ptr = make_ptr(sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
+     sfb1_ptr = make_ptr(sf_dtype, sfb1_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
+     sfb2_ptr = make_ptr(sf_dtype, sfb2_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
+     compiled_func(a_ptr, b1_ptr, sfa_ptr, sfb1_ptr, g1_ptr, m, n, k, l)
+     compiled_func(a_ptr, b2_ptr, sfa_ptr, sfb2_ptr, g2_ptr, m, n, k, l)
+     torch.nn.functional.silu(g1, inplace=True)
+     g1.mul_(g2)
+     return c
  __all__ = ["custom_kernel"]