From 2730b657c4aa92a1a342a243514bbbd68804866f Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Mon, 20 Jul 2026 10:45:56 +0800 Subject: [PATCH] [Bugfix] Fix broken NVVM caused by CuteDSL 4.6.0 (#49108) Signed-off-by: Thien Tran --- .buildkite/test_areas/kernels.yaml | 7 +++++++ vllm/cute_utils/_tcgen05.py | 8 ++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 1749b71abe..3618cb6e96 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -237,6 +237,11 @@ steps: - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py + - vllm/cute_utils/ + - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/ + - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py + - tests/kernels/mamba/test_gdn_prefill_cutedsl.py + - tests/kernels/test_bf16x3_router_gemm_cutedsl.py - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py commands: @@ -266,6 +271,8 @@ steps: - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py + - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py + - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py - pytest -v -s tests/kernels/test_ll_bf16_gemm.py # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py diff --git a/vllm/cute_utils/_tcgen05.py b/vllm/cute_utils/_tcgen05.py index 9367fa12a1..325c30498c 100644 --- a/vllm/cute_utils/_tcgen05.py +++ b/vllm/cute_utils/_tcgen05.py @@ -10,8 +10,8 @@ from cutlass.cutlass_dsl import dsl_user_op NVVM_CTA_GROUP_MAP = [ None, - nvvm.Tcgen05GroupKind.CTA_1, - nvvm.Tcgen05GroupKind.CTA_2, + nvvm.CTAGroupKind.CTA_1, + nvvm.CTAGroupKind.CTA_2, ] LDST_MAP = { "32x32b": (nvvm.Tcgen05LdStShape.SHAPE_32X32B, 1), @@ -136,7 +136,7 @@ def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None): group = NVVM_CTA_GROUP_MAP[cta_group] if cutlass.const_expr(cta_mask is not None): with cute.arch.elect_one(): - nvvm.tcgen05_commit_arrive( + nvvm.tcgen05_commit( mbar_llvm, multicast_mask=cta_mask.ir_value(loc=loc, ip=ip), group=group, @@ -145,7 +145,7 @@ def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None): ) else: with cute.arch.elect_one(): - nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip) + nvvm.tcgen05_commit(mbar_llvm, group=group, loc=loc, ip=ip) @dsl_user_op