diff --git a/pyproject.toml b/pyproject.toml index a4a4606..9366aea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,8 +18,8 @@ dependencies = [ "ninja", "pybind11", "cuda-python", - "nvidia-cutlass-dsl>=4.4.1", - "quack-kernels>=0.2.10", + "nvidia-cutlass-dsl>=4.6.0", + "quack-kernels>=0.6.1", ] [project.urls] diff --git a/python/fmha_sm100/cute/requirements.txt b/python/fmha_sm100/cute/requirements.txt index a188988..de0aa93 100644 --- a/python/fmha_sm100/cute/requirements.txt +++ b/python/fmha_sm100/cute/requirements.txt @@ -1,2 +1,2 @@ -nvidia-cutlass-dsl>=4.4.1 -quack-kernels>=0.2.10 +nvidia-cutlass-dsl>=4.6.0 +quack-kernels>=0.6.1 diff --git a/python/fmha_sm100/cute/src/common/tma_utils.py b/python/fmha_sm100/cute/src/common/tma_utils.py index 5bdc19a..2b9f8d9 100644 --- a/python/fmha_sm100/cute/src/common/tma_utils.py +++ b/python/fmha_sm100/cute/src/common/tma_utils.py @@ -120,7 +120,7 @@ def prefetch_tma_desc_raw(tma_desc_ptr, *, loc=None, ip=None): ptr_i64_align = cute_ir.assume(ptr_i64_align_ty, ptr_i64, loc=loc, ip=ip) ptr_ty = cute_ir.PtrType.get( cute_nvgpu_ir.TmaDescriptorTiledType.get(), - cute_ir.AddressSpace.gmem, + cute_ir.AddressSpace.generic, 128, ) desc_ptr = cute_ir.inttoptr(ptr_ty, ptr_i64_align, loc=loc, ip=ip) diff --git a/python/fmha_sm100/cute/src/common/utils.py b/python/fmha_sm100/cute/src/common/utils.py index 10b2b1f..98b6869 100644 --- a/python/fmha_sm100/cute/src/common/utils.py +++ b/python/fmha_sm100/cute/src/common/utils.py @@ -153,7 +153,9 @@ def make_tiled_copy_B( def mma_make_fragment_A( - smem: cute.Tensor, thr_mma: cute.core.ThrMma, swapAB: cutlass.Constexpr[bool] = False + smem: cute.Tensor, + thr_mma: cute.ThrMma, + swapAB: cutlass.Constexpr[bool] = False, ) -> cute.Tensor: if const_expr(swapAB): return mma_make_fragment_B(smem, thr_mma) @@ -162,7 +164,9 @@ def mma_make_fragment_A( def mma_make_fragment_B( - smem: cute.Tensor, thr_mma: cute.core.ThrMma, swapAB: cutlass.Constexpr[bool] = False + smem: cute.Tensor, + thr_mma: cute.ThrMma, + swapAB: cutlass.Constexpr[bool] = False, ) -> cute.Tensor: if const_expr(swapAB): return mma_make_fragment_A(smem, thr_mma) @@ -705,7 +709,14 @@ def cvt_fp4x8_e2m1_scaled_e4m3x8( from cutlass import CUDA_VERSION - if CUDA_VERSION.major > 13 or (CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 2): + # CUTLASS DSL 4.6.0 is built against CUDA 13.3, but enabling this path + # currently triggers NVVM_ERROR_COMPILATION. Keep using the original + # implementation for now. + use_native_fp4_mul = False + if use_native_fp4_mul and ( + CUDA_VERSION.major > 13 + or (CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 2) + ): out = llvm.inline_asm( llvm.StructType.get_literal([T.i32(), T.i32()]), [ diff --git a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py index f5950f3..caad594 100644 --- a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py +++ b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py @@ -1939,8 +1939,8 @@ def _wg_mma_issue( self, tiled_mma_qk: cute.TiledMma, tiled_mma_pv: cute.TiledMma, - thr0_qk: cute.core.ThrMma, - thr0_pv: cute.core.ThrMma, + thr0_qk: cute.ThrMma, + thr0_pv: cute.ThrMma, tStS: cute.Tensor, tOrP: cute.Tensor, sK: cute.Tensor, diff --git a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py index ba930a6..846df38 100644 --- a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py +++ b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py @@ -2191,8 +2191,8 @@ def _wg_mma_issue( self, tiled_mma_qk: cute.TiledMma, tiled_mma_pv: cute.TiledMma, - thr0_qk: cute.core.ThrMma, - thr0_pv: cute.core.ThrMma, + thr0_qk: cute.ThrMma, + thr0_pv: cute.ThrMma, tStS: cute.Tensor, tOrP: cute.Tensor, sK: cute.Tensor, diff --git a/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py b/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py index 8ccc2cb..7bbbc60 100644 --- a/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py +++ b/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py @@ -455,7 +455,10 @@ class SharedStorage: tmem_dealloc_mbar_ptr: Int64 tmem_holding_buf: Int32 clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size] - clc_response: cute.struct.MemRange[Int32, clc_response_size] + clc_response: cute.struct.Align[ + cute.struct.MemRange[Int32, clc_response_size], + 16, + ] sQ: cute.struct.Align[ cute.struct.MemRange[self.q_dtype, cute.cosize(sQ_layout)], self.buffer_align_bytes, @@ -726,7 +729,7 @@ def kernel( pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True) if const_expr(self.use_clc_scheduler): - clc_response_ptr = storage.clc_response.data_ptr() + clc_response_ptr = storage.clc_response.data_ptr().align(16) clc_mbar_ptr = storage.clc_mbar_ptr.data_ptr() clc_pipeline_producer_group = cutlass_pipeline.CooperativeGroup( cutlass_pipeline.Agent.Thread @@ -1536,8 +1539,8 @@ def correction_epilogue_combine( tOtO1_t2r_i = tOtO1_t2r[None, 0, 0, col_pass_idx] tOsO_r2s_i = tOsO_s2r[None, 0, 0, col_pass_idx] frg_shape = tOcO_t2r[None, 0, 0, col_pass_idx].shape - tOrO0_frg = cute.make_fragment(frg_shape, self.pv_acc_dtype) - tOrO1_frg = cute.make_fragment(frg_shape, self.pv_acc_dtype) + tOrO0_frg = cute.make_rmem_tensor(frg_shape, self.pv_acc_dtype) + tOrO1_frg = cute.make_rmem_tensor(frg_shape, self.pv_acc_dtype) is_zero_output = ( scale0 == Float32(0.0) and scale1 == Float32(0.0) ) @@ -1600,7 +1603,7 @@ def correction_rescale( frg_count: cutlass.Constexpr[int] = self.head_dim // corr_tile_size for fi in cutlass.range_constexpr(frg_count): - tOrO_frg = cute.make_fragment( + tOrO_frg = cute.make_rmem_tensor( tOrO_t2r_shape, self.pv_acc_dtype) tOtO_t2r_i = cute.make_tensor( tOtO_t2r.iterator + fi * corr_tile_size, @@ -1899,7 +1902,7 @@ def softmax_loop( tmem_load_atom_pre: cute.CopyAtom, tmem_store_atom_pre: cute.CopyAtom, tmem_store_vec_atom_pre: cute.CopyAtom, - thr_mma_qk_pre: cute.core.ThrMma, + thr_mma_qk_pre: cute.ThrMma, pipeline_s_p_o: pipeline.PipelineAsync, pipeline_p_lastsplit: pipeline.PipelineAsync, pipeline_sm_stats: pipeline.PipelineAsync, diff --git a/python/fmha_sm100/cute/src/sm100/fwd_decode/tile_scheduler.py b/python/fmha_sm100/cute/src/sm100/fwd_decode/tile_scheduler.py index 599dace..6f0b3fb 100644 --- a/python/fmha_sm100/cute/src/sm100/fwd_decode/tile_scheduler.py +++ b/python/fmha_sm100/cute/src/sm100/fwd_decode/tile_scheduler.py @@ -180,7 +180,9 @@ def _clc_response_to_work( # CLC responses are 16B opaque records. The scheduler warp can query # the next stage before all consumer warps have read the current one, # so each pipeline stage needs its own response slot. - response_ptr = self._clc_response_ptr + response_stage * Int32(4) + response_ptr = ( + self._clc_response_ptr + response_stage * Int32(4) + ).align(16) m_idx, n_idx, l_idx, is_valid = cute.arch.clc_response( response_ptr, loc=loc, ip=ip) cute.arch.fence_proxy("async.shared", space="cta") @@ -239,7 +241,9 @@ def advance_to_next_work( ): if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): assert mbarrier_addr is not None - response_ptr = self._clc_response_ptr + response_stage * Int32(4) + response_ptr = ( + self._clc_response_ptr + response_stage * Int32(4) + ).align(16) with cute.arch.elect_one(): cute.arch.issue_clc_query( mbarrier_addr, response_ptr, loc=loc, ip=ip) diff --git a/requirements.txt b/requirements.txt index 1575094..416c38a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,5 +10,5 @@ pybind11 cuda-python # CuTe-DSL sparse attention backend -nvidia-cutlass-dsl>=4.4.1 -quack-kernels>=0.2.10 +nvidia-cutlass-dsl>=4.6.0 +quack-kernels>=0.6.1