Skip to content

Commit e15ae4e

Browse files
committed
fix(cute_dsl/attention): use string literals for fence_proxy (cu13 wheel compat)
Fixes #3071. `cute.arch.fence_proxy` is documented as taking string literals (``"async.shared"``, ``"cta"``); the wrapper does ``ProxyKind.from_str`` / ``SharedSpace.from_str`` internally. The enum re-exports ``cute.arch.ProxyKind`` and ``cute.arch.SharedSpace`` are gated upstream behind ``cutlass_dsl.target_version(exact_version="12.9")`` and so are absent from the cu13 wheel that the ``flashinfer-ci-cu130`` docker image uses, producing: AttributeError: module 'cutlass.cute.arch' has no attribute 'ProxyKind' at correction.py:267 and softmax.py:448. On cu12 wheels the enum form still works but emits ``DeprecationWarning: Passing enum member directly to SharedSpace.from_str() is deprecated. Please use string literals instead`` (visible in our local test runs). Switching the two call sites to the string form fixes both the cu13 breakage and the cu12 deprecation warning, and is the documented stable API on every flashinfer-supported cutlass-dsl 4.4.x wheel. Verified: python -m pytest tests/attention/test_modular_fmha_prefill.py \ tests/attention/test_cute_dsl_mla_decode.py → 456 passed, 20 skipped (no SharedSpace/ProxyKind warnings emitted). Made-with: Cursor
1 parent 6e67879 commit e15ae4e

File tree

2 files changed

+2
-8
lines changed

2 files changed

+2
-8
lines changed

flashinfer/cute_dsl/attention/roles/correction.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,10 +263,7 @@ def epilog(
263263
cute.copy(tiled_smem_store, tSMrO, tTMEM_LOADsO_i)
264264

265265
# fence view async shared
266-
cute.arch.fence_proxy(
267-
cute.arch.ProxyKind.async_shared,
268-
space=cute.arch.SharedSpace.shared_cta,
269-
)
266+
cute.arch.fence_proxy("async.shared", space="cta")
270267

271268
@cute.jit
272269
def run(

flashinfer/cute_dsl/attention/roles/softmax.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -444,10 +444,7 @@ def softmax_epilog(
444444
tSMrO.store(o_vec.to(self.o_dtype))
445445
cute.copy(tiled_smem_store, tSMrO, tTMEM_LOADsO_i)
446446

447-
cute.arch.fence_proxy(
448-
cute.arch.ProxyKind.async_shared,
449-
space=cute.arch.SharedSpace.shared_cta,
450-
)
447+
cute.arch.fence_proxy("async.shared", space="cta")
451448

452449
# For both softmax0 and softmax1 warp group
453450
@cute.jit

0 commit comments

Comments
 (0)