[CuTe DSL] Fix FP8 MLA persistent perf regression and ProxyKind cu13 wheel breakage (#3132)

## Summary

Two functional fixes plus two defensive refactors on the modular CuTe
DSL attention kernels added in #2805:

1. **`perf(cute_dsl/mla)`** — recover a ~6.4 % geomean perf regression
on FP8 MLA decode persistent (`is_var_seq=False`) versus trtllm-gen on
B200, by unblocking the cute-DSL JIT's constant folding around the MMA
dispatch loop. Cute-DSL now beats trtllm-gen on 9/10 of the standard
`seq_len=8192, num_heads=128` configs (was 0/10).
2. **`fix(cute_dsl/attention)`** — closes
[#3071](#3071): the
`cu13` cutlass-dsl wheel doesn't re-export `cute.arch.ProxyKind` /
`cute.arch.SharedSpace` (gated upstream behind
`cutlass_dsl.target_version("12.9")`), causing `AttributeError` at
`correction.py:267` and `softmax.py:448` on the CI's
`flashinfer-ci-cu130` image. Switching to the documented string-literal
API (`fence_proxy("async.shared", space="cta")`) works on both cu12.9
and cu13 wheels and silences the existing `Passing enum member directly
to SharedSpace.from_str() is deprecated` warning on cu12.9.
3. **`refactor(cute_dsl/mla)`** + **`refactor(cute_dsl/fmha)`** — apply
the same fresh-local-`TiledMma` + `range_constexpr` pattern that the FP8
perf fix introduced to the BF16/FP16 MLA decode and FMHA prefill MMA
roles. Wall-clock perf is unchanged for these (the back-end was already
collapsing the buggy form to identical SASS for these dtypes/shapes),
but the IR is cleaner and the helpers' API is consistent across all
three MMA roles, protecting against future cute-DSL back-end heuristic
changes that could regress one variant the way FP8 did.

## Commits

| SHA | What |
|---|---|
| `ba59eb61` | `perf(cute_dsl/mla): unblock MMA-warp constant folding
for FP8 MLA decode (persistent)` |
| `4e66703c` | `refactor(cute_dsl/mla): isolate MMA-warp TiledMma
mutations in BF16/FP16 helpers` |
| `6e67879d` | `refactor(cute_dsl/fmha): isolate MMA-warp TiledMma
mutations in prefill helpers` |
| `e15ae4ea` | `fix(cute_dsl/attention): use string literals for
fence_proxy (cu13 wheel compat)` |

## Root cause analysis (FP8 perf regression)

PR #2805 introduced two patterns in the FP8 MMA role that prevented the
cute-DSL JIT from emitting an unrolled MMA dispatch loop with
compile-time-constant `ACCUMULATE` bits:

1. The inner k-block loops in `_gemm_qk_latent_one_stage`,
`_gemm_qk_rope_one_stage`, and `_gemm_pv_one_stage` had been switched
from `cutlass.range_constexpr` to `cutlass.range`. The original PR's
comment explained this as a workaround for an SSA-dominance failure:
when the helper's unrolled `tiled_mma.set(ACCUMULATE, ...)` chain
inlined into the persistent `while` loop, the chain rooted in the
caller's `TiledMma` SSA value and the loop's back-edge couldn't pick a
dominating value to thread back. The runtime-loop workaround compiles
but loses tcgen05 dispatch throughput. Fixed by having each helper
construct its own local `TiledMma` via
`sm100_utils.make_trivial_tiled_mma(...)` and mutate that local instance
only — the unrolled chain now dies inside the helper frame and
`range_constexpr` is safe again. Same pattern the compute role already
uses.

2. PV per-tile "first-iteration overwrite, then accumulate" state was
tracked in a Python `bool` (`pv_accumulated`) reassigned inside the
dynamic k_tile loop. The cute-DSL JIT demoted that bool to a runtime
`i1` carried through the loop, making `accumulate or p_stage > 0` a
runtime OR and forcing every PV MMA to compute its `ACCUMULATE` bit at
runtime. Fixed by storing the flag on `tiled_mma_pv.set(ACCUMULATE,
...)` / `.get(ACCUMULATE)` so the field becomes type-side metadata that
the JIT propagates as a per-iteration compile-time constant — same
pattern the BF16 `mla_mma.py` already used.

## Verification

### FP8 MLA persistent perf (Table 1 from PR #2743's reproducer, B200,
median of 3 runs, CUPTI + CUDA graph + cold L2)

| Config | trt | cute_base | cute_fix | sp_base | sp_fix | cute Δ% |
| :----------- | -----: | --------: | -------: | ------: | -----: |
-------: |
| B= 1, q=1 | 0.0157 | 0.0158 | 0.0161 | 1.00x | 0.99x | +1.9 % |
| B= 32, q=1 | 0.0521 | 0.0547 | 0.0522 | 0.96x | 1.00x | −4.7 % |
| B= 64, q=1 | 0.0784 | 0.0795 | 0.0734 | 0.99x | 1.06x | −7.7 % |
| B=128, q=1 | 0.1448 | 0.1494 | 0.1391 | 0.97x | 1.04x | −6.9 % |
| B=256, q=1 | 0.2890 | 0.3020 | 0.2808 | 0.94x | 1.03x | −7.0 % |
| B= 1, q=4 | 0.0190 | 0.0189 | 0.0181 | 1.00x | 1.05x | −3.8 % |
| B= 32, q=4 | 0.1314 | 0.1447 | 0.1257 | 0.91x | 1.05x | −13.2 % |
| B= 64, q=4 | 0.2627 | 0.2810 | 0.2540 | 0.93x | 1.03x | −9.6 % |
| B=128, q=4 | 0.4905 | 0.5069 | 0.4753 | 0.95x | 1.03x | −6.2 % |
| B=256, q=4 | 1.0050 | 1.0456 | 0.9869 | 0.95x | 1.02x | −5.6 % |
| **geomean** | | | | **0.958x** | **1.029x** | **−6.4 %** |

Geomean speedup vs trtllm-gen flipped from **0.958x → 1.029x**. cute-dsl
now beats trtllm-gen on 9/10 configs.

### IR / SASS evidence (FP8 MLA persistent kernel)

| Metric | Buggy | Fixed |
| :--- | ---: | ---: |
| cubin size | 156 296 B | **126 648 B (−19 %)** |
| `UTCQMMA` (FP8 MMA dispatch) | 103 | **52 (−50 %)** |
| `BSSY` / `BSYNC` (structured-sync regions) | 88 / 88 | **6 / 6 (−93
%)** |
| `VOTEU` (warp predicate votes) | 241 | **1 (−99 %)** |
| `WARPSYNC` | 67 | **2 (−97 %)** |
| `scf.while` outer carry: TiledMma type | 0 | **1** (PV TiledMma
replaces the demoted `i1`) |
| `scf.while` outer carry: `i1` (Python bool) | 2 | **0** |
| `scf.while` outer carry: `i32` | 31 | **20** (intermediate state from
inner runtime loops gone) |

### Other paths (defensive refactors — unchanged perf, same as
predicted)

- BF16 MLA persistent: cute geomean Δ +0.12 % (single-run noise, SASS
bit-identical)
- FP8 MLA var-seq: cute geomean Δ +0.52 % (noise)
- BF16 MLA var-seq: cute geomean Δ +0.79 % (noise)
- BF16 FMHA prefill (`benchmarks/bench_blackwell_attention_cutedsl.py`):
cute geomean Δ −0.31 % (noise)

### ProxyKind / cu13 wheel (verified on both wheel variants)

| Wheel | `target_version("12.9")` | `cute.arch.ProxyKind` | Pre-fix
code | Post-fix code |
|---|:---:|:---:|:---:|:---:|
| **cu12.9** (`libs-base==4.4.2`) | True | exists (with deprecation
warning) | passes (warning) | **passes** |
| **cu13** (`libs-base + libs-cu13==4.4.2`) | False | **AttributeError**
| **fails** (#3071, exact bit-identical reproduction) | **passes** |

## Test plan

- [x] `tests/attention/test_cute_dsl_mla_decode.py` — 297 passed
(cu12.9), includes 24 FP8-specific cases
- [x] `tests/attention/test_modular_fmha_prefill.py` — 159 passed, 20
skipped (cu12.9)
- [x] Same test suites on cu13 wheel — 456 passed, 20 skipped (closes
#3071 cleanly)
- [x] `bench_pr2743_reproduce.py` (Tables 1-4, 3 runs each for FP8
fixed-len) — see numbers above
- [x] `benchmarks/bench_blackwell_attention_cutedsl.py` — perf flat as
expected
- [x] Pre-commit (mypy, ruff, etc.) — passes for all 4 commits

Made with [Cursor](https://cursor.com)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Attention path now explicitly configures operand data types at runtime
for more consistent numeric behavior (including FP8 paths).
* Matrix-multiply helpers now build local compute instances, reducing
external coupling and simplifying orchestration.
  * PV accumulation logic streamlined for clearer accumulate semantics.
* **Bug Fix**
* Tightened post-shared-memory synchronization to improve stability and
determinism.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: pgera

flashinfer/cute_dsl/attention/mla_decode.py

@@ -188,6 +188,7 @@ def _reinterpret_3d_kv(t):
         self.pt_loader_role = MLAPageTableLoaderRole(self.config)
         self.loader_role = MLALoaderRole(self.config)
         self.mma_role = MLAMmaRole(self.config, self.mainloop)
+        self.mma_role.set_dtypes(self.q_dtype, self.v_dtype)
         self.compute_role = MLAComputeRole(self.config, fusion=self.fusion)
         self.compute_role.set_dtypes(self.q_dtype)
         self.compute_role.set_barriers(self.softmax_exchange_sync_bar)

flashinfer/cute_dsl/attention/mla_decode_fp8.py

@@ -191,6 +191,7 @@ def _reinterpret_3d_kv(t):
         self.loader_k_role = MLAFP8LoaderKRole(self.config)
         self.loader_v_role = MLAFP8LoaderVRole(self.config)
         self.mma_role = MLAMmaFP8Role(self.config, self.mainloop)
+        self.mma_role.set_dtypes(self.q_dtype, self.v_dtype)
         self.compute_role = MLAComputeRole(self.config, fusion=self.fusion)
         self.compute_role.set_dtypes(self.q_dtype)
         self.compute_role.set_barriers(self.softmax_exchange_sync_bar)

flashinfer/cute_dsl/attention/prefill.py

@@ -206,6 +206,13 @@ def __call__(
             threads_per_warp=self.schedule.threads_per_warp,
             has_logits_transform=self.has_logits_transform,
         )
+        self.mma_role.set_dtypes(
+            self.q_dtype,
+            self.v_dtype,
+            self.q_major_mode,
+            self.k_major_mode,
+            self.v_major_mode,
+        )
         self.softmax_role.set_dtypes(self.q_dtype, self.o_dtype)
 
         lp = build_fmha_launch_params(
@@ -495,7 +502,6 @@ def kernel(
         if warp_idx == self.schedule.mma_warp_id:
             cute.arch.warpgroup_reg_dealloc(self.schedule.num_regs_other)
             self.mma_role.run(
-                qk_tiled_mma,
                 pv_tiled_mma,
                 tStS0,
                 tStS1,

flashinfer/cute_dsl/attention/roles/correction.py

@@ -263,10 +263,7 @@ def epilog(
             cute.copy(tiled_smem_store, tSMrO, tTMEM_LOADsO_i)
 
         # fence view async shared
-        cute.arch.fence_proxy(
-            cute.arch.ProxyKind.async_shared,
-            space=cute.arch.SharedSpace.shared_cta,
-        )
+        cute.arch.fence_proxy("async.shared", space="cta")
 
     @cute.jit
     def run(

flashinfer/cute_dsl/attention/roles/mla_mma.py

@@ -19,11 +19,22 @@
 each helper as ``k_block != 0 or accumulate``, making it deterministic from
 the parameter and inner loop index. Callers compute the parameter from their
 own loop position — never from ``tiled_mma.get()`` after a sub-method return.
+
+Inner k-block loops use ``cutlass.range_constexpr`` (compile-time unrolled)
+for maximum tcgen05 MMA dispatch throughput.  Each GEMM helper constructs
+its own local ``TiledMma`` via ``make_trivial_tiled_mma`` so that the
+``.set(ACCUMULATE, ...)`` mutations stay inside the helper's frame and
+never leak SSA values across the persistent tile-scheduler ``while`` loop
+in ``run()``.  Same isolation pattern the compute role and the FP8 MMA
+role (mla_mma_fp8.py) already use.
 """
 
+from typing import Type
+
 import cutlass
 import cutlass.cute as cute
 import cutlass.cute.nvgpu.tcgen05 as tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.pipeline import PipelineProducer, PipelineConsumer
 from types import SimpleNamespace
 
@@ -56,6 +67,56 @@ def __init__(self, config: MLAConfig, mainloop: MLAMainloopSpec):
         self.iterations_pv_n = config.iterations_pv_n
         self.enable_pdl = config.enable_pdl
         self.is_var_split_kv = config.is_var_split_kv
+        self.use_2cta_instrs = config.use_2cta_instrs
+        self.acc_dtype = config.acc_dtype
+        # self.q_dtype and self.v_dtype are populated by set_dtypes() — the
+        # operand element types are only known at __call__ time on the kernel.
+
+    def set_dtypes(
+        self,
+        q_dtype: Type[cutlass.Numeric],
+        v_dtype: Type[cutlass.Numeric],
+    ) -> None:
+        """Set tensor element types discovered at call time.
+
+        Required so the GEMM helpers can reconstruct local TiledMma
+        instances via ``make_trivial_tiled_mma``.
+        """
+        self.q_dtype: Type[cutlass.Numeric] = q_dtype
+        self.v_dtype: Type[cutlass.Numeric] = v_dtype
+
+    @cute.jit
+    def _make_local_qk_mma(self) -> cute.TiledMma:
+        """Fresh QK TiledMma — mutations on this instance never escape the
+        helper that constructs it, so the inner k-block loop can use
+        ``range_constexpr`` without leaking SSA values into the enclosing
+        persistent ``while`` loop in ``run()``."""
+        cta_group = (
+            tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
+        )
+        return sm100_utils.make_trivial_tiled_mma(
+            self.q_dtype,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.K,
+            self.acc_dtype,
+            cta_group,
+            self.mma_qk_tiler[:2],
+        )
+
+    @cute.jit
+    def _make_local_pv_mma(self) -> cute.TiledMma:
+        """Fresh PV TiledMma — same isolation rationale as ``_make_local_qk_mma``."""
+        cta_group = (
+            tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
+        )
+        return sm100_utils.make_trivial_tiled_mma(
+            self.v_dtype,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.MN,
+            self.acc_dtype,
+            cta_group,
+            self.mma_pv_tiler[:2],
+        )
 
     # ------------------------------------------------------------------
     #  Tile count
@@ -97,30 +158,32 @@ def _get_k_tile_count(
     #  state back via TiledMma mutations (they would be invisible to the
     #  caller due to SSA pass-by-value at the @cute.jit boundary).
     #
-    #  Inner k-block loops use ``cutlass.range()`` (dynamic scf.for),
-    #  NOT ``cutlass.range_constexpr()`` (compile-time unroll).
-    #  range_constexpr unrolls tiled_mma.set() calls into the enclosing
-    #  scope, producing SSA values that leak across dynamic while-loop
-    #  yields. range() keeps the .set() inside an scf.for scope where
-    #  SSA carry-through is handled correctly.
+    #  Inner k-block loops use ``cutlass.range_constexpr`` (compile-time
+    #  unrolled) for maximum tcgen05 MMA dispatch throughput.  To prevent
+    #  the unrolled ``tiled_mma.set(ACCUMULATE, ...)`` mutations from
+    #  leaking SSA values into the enclosing persistent ``while`` loop
+    #  in ``run()`` (which would cause SSA-dominance failures), each
+    #  helper constructs a fresh local TiledMma via
+    #  ``make_trivial_tiled_mma`` and mutates that local instance only.
+    #  The caller's TiledMma is never touched by the helper.
     # ------------------------------------------------------------------
 
     @cute.jit
     def _gemm_qk_latent_one_stage(
         self,
         qk_params: SimpleNamespace,
-        tiled_mma_qk: cute.TiledMma,
         s_stage_index: cutlass.Int32,
         kv_stage_index: cutlass.Int32,
         q_stage: int,
         accumulate: bool,
     ):
         """Compute one QK-latent stage: inner k-block GEMM loop."""
+        local_mma = self._make_local_qk_mma()
         tStS = qk_params.tStS_staged[None, None, None, s_stage_index]
-        for k_block in cutlass.range(cute.size(qk_params.tSrQ.shape[2])):
-            tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, k_block != 0 or accumulate)
+        for k_block in cutlass.range_constexpr(cute.size(qk_params.tSrQ.shape[2])):
+            local_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0 or accumulate)
             cute.gemm(
-                tiled_mma_qk,
+                local_mma,
                 tStS,
                 qk_params.tSrQ[None, None, k_block, q_stage],
                 qk_params.tSrKC[None, None, k_block, kv_stage_index],
@@ -131,18 +194,18 @@ def _gemm_qk_latent_one_stage(
     def _gemm_qk_rope_one_stage(
         self,
         qk_params: SimpleNamespace,
-        tiled_mma_qk: cute.TiledMma,
         s_stage_index: cutlass.Int32,
         kv_stage_index: cutlass.Int32,
         q_stage: int,
         accumulate: bool,
     ):
         """Compute one QK-rope stage: inner k-block GEMM loop."""
+        local_mma = self._make_local_qk_mma()
         tStS = qk_params.tStS_staged[None, None, None, s_stage_index]
-        for k_block in cutlass.range(self.rope_dim // tiled_mma_qk.shape_mnk[2]):
-            tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, k_block != 0 or accumulate)
+        for k_block in cutlass.range_constexpr(self.rope_dim // local_mma.shape_mnk[2]):
+            local_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0 or accumulate)
             cute.gemm(
-                tiled_mma_qk,
+                local_mma,
                 tStS,
                 qk_params.tSrQ_rope[None, None, k_block, q_stage],
                 qk_params.tSrKC[None, None, k_block, kv_stage_index],
@@ -153,19 +216,19 @@ def _gemm_qk_rope_one_stage(
     def _gemm_pv_one_stage(
         self,
         pv_params: SimpleNamespace,
-        tiled_mma_pv: cute.TiledMma,
         p_stage_index: cutlass.Int32,
         kv_stage_index: cutlass.Int32,
         p_stage: int,
         acc_stage: int,
         accumulate: bool,
     ):
         """Compute one PV stage: inner k-block GEMM loop."""
+        local_mma = self._make_local_pv_mma()
         tOtO = pv_params.tOtO_staged[None, None, None, acc_stage]
-        for k_block in cutlass.range(pv_params.tOrP.shape[2]):
-            tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, k_block != 0 or accumulate)
+        for k_block in cutlass.range_constexpr(pv_params.tOrP.shape[2]):
+            local_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0 or accumulate)
             cute.gemm(
-                tiled_mma_pv,
+                local_mma,
                 tOtO,
                 pv_params.tOrP[
                     None,
@@ -284,7 +347,6 @@ def run(
                         kv_handle = load_kv_consumer.wait_and_advance()
                         self._gemm_qk_latent_one_stage(
                             mma_qk_params,
-                            tiled_mma_qk,
                             s_handle.index,
                             kv_handle.index,
                             q_stage,
@@ -295,7 +357,6 @@ def run(
                         kv_handle = load_kv_consumer.wait_and_advance()
                         self._gemm_qk_rope_one_stage(
                             mma_qk_params,
-                            tiled_mma_qk,
                             s_handle.index,
                             kv_handle.index,
                             q_stage,
@@ -313,7 +374,6 @@ def run(
                             kv_handle = load_kv_consumer.wait_and_advance()
                             self._gemm_qk_latent_one_stage(
                                 mma_qk_params,
-                                tiled_mma_qk,
                                 s_handle.index,
                                 kv_handle.index,
                                 q_stage,
@@ -324,7 +384,6 @@ def run(
                             kv_handle = load_kv_consumer.wait_and_advance()
                             self._gemm_qk_rope_one_stage(
                                 mma_qk_params,
-                                tiled_mma_qk,
                                 s_handle.index,
                                 kv_handle.index,
                                 q_stage,
@@ -343,7 +402,6 @@ def run(
                                 kv_handle = load_kv_consumer.wait_and_advance()
                                 self._gemm_pv_one_stage(
                                     mma_pv_params,
-                                    tiled_mma_pv,
                                     p_handle.index,
                                     kv_handle.index,
                                     p_stage,
@@ -368,7 +426,6 @@ def run(
                             kv_handle = load_kv_consumer.wait_and_advance()
                             self._gemm_pv_one_stage(
                                 mma_pv_params,
-                                tiled_mma_pv,
                                 p_handle.index,
                                 kv_handle.index,
                                 p_stage,