Skip to content

Fix CUTLASS DSL 4.6.0 compatibility #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zyongye merged 1 commit into from
Jul 20, 2026
Merged

Fix CUTLASS DSL 4.6.0 compatibility #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
4 changes: 2 additions & 2 deletions python/fmha_sm100/cute/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion python/fmha_sm100/cute/src/common/tma_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
17 changes: 14 additions & 3 deletions python/fmha_sm100/cute/src/common/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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()]),
[
Expand Down
4 changes: 2 additions & 2 deletions python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
15 changes: 9 additions & 6 deletions python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 6 additions & 2 deletions python/fmha_sm100/cute/src/sm100/fwd_decode/tile_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@arpera arpera Jul 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently vllm requires quack-kernels>=0.4.0. In that case do we need to align vllm requirements with MSA requirements?

@lengrongfu lengrongfu Jul 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can in this pr vllm-project/vllm#49016 to together update.

@wzhao18 wzhao18 Jul 19, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@arpera Should we pin to an exact version? With quack-kernels>=xxx I believe it will just install the latest version.

@lengrongfu Sounds good. We can update the version there when we update the MSA commit hash. Thanks!

@arpera arpera Jul 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there is no need to pin an exact version here since we only need at least 0.6.1. Also, leave please a comment in the files where you update version of quack-kernels to know exactly why we need at least 0.6.1 installed.