MXFP8 Perf Improvement

python/cudnn/sdpa/fwd/api_dsl.py

@@ -9,8 +9,9 @@
 import math
 import os
 from abc import abstractmethod
+from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Callable, Hashable, Optional
+from typing import Callable, Hashable, Iterator, Optional
 
 import torch
 from cuda.bindings import driver as cuda
@@ -26,6 +27,7 @@
     MASK_NONE,
     MASK_PADDED,
     MASK_SWA,
+    SCHED_LPT,
     SCHED_NATURAL,
 )
 from cudnn.sdpa.fwd.config_sm100 import TemplateParams as Sm100TemplateParams
@@ -73,6 +75,25 @@
 _WS_ALIGN = 128
 
 
+@contextmanager
+def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch.device) -> Iterator[None]:
+    """Run PyTorch work on the CUDA stream used for the kernel launch."""
+    if current_stream is None:
+        yield
+        return
+    handle = int(current_stream)
+    torch_current = torch.cuda.current_stream(device)
+    torch_default = torch.cuda.default_stream(device)
+    if handle == torch_current.cuda_stream:
+        launch_stream = torch_current
+    elif handle == torch_default.cuda_stream:
+        launch_stream = torch_default
+    else:
+        launch_stream = torch.cuda.ExternalStream(handle, device=device)
+    with torch.cuda.stream(launch_stream):
+        yield
+
+
 def ws_align(nbytes: int) -> int:
     """Round a scratch-chunk size up to the carve alignment (128 B)."""
     return -(-int(nbytes) // _WS_ALIGN) * _WS_ALIGN
@@ -566,6 +587,9 @@ def compile(self) -> None:
         # picks the fused path with no user action.
         mxfp8 = self._fp8 and not self._pertensor
         fused_ldtm_stat = mxfp8 and (self._device_cc == (10, 3))
+        sched_policy = self.sched_policy
+        if mxfp8 and sched_policy == SCHED_NATURAL and (self.mask_flags & MASK_CAUSAL):
+            sched_policy = SCHED_LPT
         params = Sm100TemplateParams(
             dtype_qkv=_SM100_DTYPE_QKV_CODE[self.dtype],
             dtype_o=_SM100_DTYPE_QKV_CODE[self.dtype_o],
@@ -575,7 +599,7 @@ def compile(self) -> None:
             has_sink=self.has_sink,
             seq_kv_lens_present=self.seq_kv_lens_present,
             seq_q_lens_present=self.seq_q_lens_present,
-            sched_policy=self.sched_policy,
+            sched_policy=sched_policy,
             thd_varlen=self.thd,
             fused_ldtm_stat=fused_ldtm_stat,
         )
@@ -971,6 +995,17 @@ def _execute_mxfp8(
         )
         o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device))
 
+        amax_o_buf = (
+            amax_o.reshape(-1)[:1]
+            if amax_o is not None
+            else self._dummy("amax_o", device, lambda: torch.zeros(1, dtype=torch.float32, device=device))
+        )
+        # Must be enqueued on the SAME stream as the kernel launch below, else the
+        # reset and the kernel's atomicMax are unordered (and the reset is missing
+        # from a CUDA-graph capture taken on the handle's stream).
+        with _torch_stream_context(current_stream, device):
+            amax_o_buf.zero_()
+
         self._compiled_kernel(
             Q,
             K,
@@ -980,6 +1015,7 @@ def _execute_mxfp8(
             sf_k_v,
             sf_v_v,
             lse,
+            amax_o_buf,
             sinks_t,
             seq_kv_t,
             o_desc_dummy,
@@ -992,12 +1028,6 @@ def _execute_mxfp8(
         )
         if o_needs_copy_back:
             O_view.copy_(O)
-        if amax_o is not None:
-            # Cast to fp32 before abs(): torch has no abs for fp8 output dtypes. For
-            # FP8 O this is the amax of the (already-quantized) output — exact when
-            # no saturation, which the direct-cast (scale=1) epilogue guarantees for
-            # in-range attention outputs.
-            amax_o.reshape(-1)[:1] = O_view.to(torch.float32).abs().max()
         self._logger.debug("execute (MXFP8) completed")
 
     def _execute_fp8(
@@ -1066,8 +1096,11 @@ def _scalar(t, default=1.0):
         # dividing by scale_o below yields the pre-quant output amax.
         amax_s_buf = amax_s.reshape(-1)[:1] if amax_s is not None else self._dummy("amax_s", device, lambda: torch.zeros(1, dtype=torch.float32, device=device))
         amax_o_buf = amax_o.reshape(-1)[:1] if amax_o is not None else self._dummy("amax_o", device, lambda: torch.zeros(1, dtype=torch.float32, device=device))
-        amax_s_buf.zero_()
-        amax_o_buf.zero_()
+        # Same-stream ordering as MXFP8: the resets must precede the kernel's
+        # atomicMax on the launch stream, not on torch's current stream.
+        with _torch_stream_context(current_stream, device):
+            amax_s_buf.zero_()
+            amax_o_buf.zero_()
 
         self._compiled_kernel(
             Q,

python/cudnn/sdpa/fwd/config_sm100.py

@@ -32,6 +32,7 @@
     MASK_NONE,
     MASK_PADDED,
     MASK_SWA,
+    SCHED_LPT,
     SCHED_NATURAL,
 )
 
@@ -108,8 +109,8 @@ def _validate_params(flavor: str, k: TemplateParams) -> None:
             raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT is dense-only (THD carries per-sequence Q lengths via cu_seqlens)")
         if not k.seq_kv_lens_present:
             raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT requires SEQ_KV_LENS_PRESENT (padding mask)")
-    if k.sched_policy != SCHED_NATURAL:
-        raise ValueError(f"{flavor}: only SCHED_NATURAL (0) is wired up; got {k.sched_policy}")
+    if k.sched_policy not in (SCHED_NATURAL, SCHED_LPT):
+        raise ValueError(f"{flavor}: only SCHED_NATURAL (0) / SCHED_LPT (1) are wired up; got {k.sched_policy}")
 
 
 # ---------------------------------------------------------------------------

python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py

@@ -22,10 +22,7 @@
   3. **Manual row-max** on the MASK_NONE fast path — ``tcgen05_ld`` +
      ``row_max_reduction`` (the masked path's pattern) instead of
      ``tmem_load_max_reduction_tile`` (LDTM.STAT).
-  4. SM100 launch: ``cluster=(CTA_MMA, 1, 1)``; compile drops
-     ``--ptxas-options -uumn`` (cfence not honored on SM100) and the
-     ``cute.nvgpu`` scheduling hints (cfence / sched_res_busy) that
-     ``-uumn`` gated.
+  4. SM100 launch: ``cluster=(CTA_MMA, 1, 1)``.
 
 Supported: FP16 / BF16 (``DTYPE_QKV ∈ {2, 3}``); masks none / causal / SWA /
 padded and all pairwise combos (causal+swa, causal+padded, swa+padded);
@@ -1730,8 +1727,6 @@ def _correction_warp_group(
                 # vec_scale_pair emits nvvm.mul_packed_f32x2 → FMUL2.  Without it
                 # plain o_chunk*alpha lowers to scalar FMUL inside the runtime-if
                 # (downstream fp32 tcgen05_st doesn't force packed regs).
-                # Per-chunk cfence keeps ptxas from sinking the next chunk's
-                # tcgen05_ld ahead of the prior tcgen05_st (TMEM-col anti-dep).
                 if ~all_alpha_one:
                     for chunk_idx in cutlass.range_constexpr(N_CHUNKS_O):
                         o_addr = tmem_base_iter + cutlass.Int32(tmem_O_off + chunk_idx * O_CHUNK)
@@ -2123,7 +2118,5 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128,
         cutlass.Int32(0),
         fake_seq_q_lens,
         stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False),
-        # SM100: drop --ptxas-options -uumn — cfence is not honored on
-        # Blackwell, and this kernel emits no cfence.
         options="--enable-tvm-ffi",
     )

python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py

@@ -1235,7 +1235,7 @@ def _softmax_kv_body(
     reg_S = RegTile(reg_S_vec, size=CFG.TILE_N)
     current_max = current_max_unscaled * scale_log2
 
-    # cfence pins ptxas scheduler from hoisting stat-store across the wg sync.
+    # sync the warpgroups before the stat-store.
     if sub_tile_id == 1:
         nvvm.barrier_cta_sync(barrier_id=8, thread_count=256)
 
@@ -2043,5 +2043,5 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128)
         fake_amax_s,
         fake_amax_o,
         stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False),
-        options="--enable-tvm-ffi",  # SM100: cfence not honored / not emitted
+        options="--enable-tvm-ffi",
     )

python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py

@@ -32,6 +32,7 @@
 from cutlass._mlir.dialects import arith
 
 import cutlass
+from cutlass.base_dsl.typing import Pointer
 from cutlass.experimental import primitives as prims
 import cutlass.cute as cute
 import cuda.bindings.driver as _cuda_driver  # noqa: F401
@@ -330,6 +331,7 @@ def _kernel(
     tma_k_sf_desc: cutlass.GridConstant[tmap.TensorMap],
     tma_v_sf_desc: cutlass.GridConstant[tmap.TensorMap],
     lse_tensor: cute.Tensor,
+    amax_o_tensor: cute.Tensor,
     sinks_tensor: cute.Tensor,
     seq_kv_lens_tensor: cute.Tensor,
     o_desc_words: cute.Tensor,
@@ -573,6 +575,7 @@ def _kernel(
             bars=bars,
             sched=sched,
             lse_tensor=lse_tensor,
+            amax_o_tensor=amax_o_tensor,
             sinks_tensor=sinks_tensor,
             seq_kv_lens_tensor=seq_kv_lens_tensor,
             n_q_supers=n_q_supers,
@@ -1712,7 +1715,7 @@ def _softmax_kv_body(
 
     current_max = cute.math.max(max_a, max_b) * scale_log2
 
-    # cfence pins ptxas from hoisting stat-store work across the wg0/wg1 sync.
+    # sync the two softmax warpgroups before the stat-store.
     if sub_tile_id == 1:
         nvvm.barrier_cta_sync(barrier_id=8, thread_count=256)
 
@@ -1741,7 +1744,6 @@ def _softmax_kv_body(
     nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.STORE)
     bars.mb_stat_full[sub_tile_id].arrive()
 
-    # SchedResBusyXU64 around exp2 biases ptxas to keep MUFU/exp pipe busy across dependent FFMA chain.
     reg_S_a = reg_S_a * scale_log2 - new_total_max
     reg_S_b = reg_S_b * scale_log2 - new_total_max
     reg_P_a = cute.math.exp2(reg_S_a, fastmath=True)
@@ -1968,6 +1970,7 @@ def _correction_warp_group(
     bars,
     sched,
     lse_tensor: cute.Tensor,
+    amax_o_tensor: cute.Tensor,
     sinks_tensor: cute.Tensor,
     seq_kv_lens_tensor,
     n_q_supers,
@@ -2056,7 +2059,6 @@ def _correction_warp_group(
                 bars.mb_bmm2_done[qs].wait(bmm2_done_phase)
 
                 # vec_scale_pair emits mul_packed_f32x2; without it the DSL lowers to scalar FMUL inside this runtime-if.
-                # cfence after each chunk keeps ptxas from sinking next tcgen05_ld ahead of prior tcgen05_st (TMEM anti-dep).
                 if ~all_alpha_one:
                     for chunk_idx in cutlass.range_constexpr(N_CHUNKS_O):
                         o_addr = tmem_base_iter + cutlass.Int32(tmem_O_off + chunk_idx * O_CHUNK)
@@ -2126,18 +2128,23 @@ def _correction_warp_group(
                 _cu = cutlass.make_array_view(seq_kv_lens_tensor)
                 _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx])
                 _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b
-                if q_row_global < _s_q_b:
+                _row_valid = q_row_global < _s_q_b
+                if _row_valid:
                     lse_arr = cutlass.make_array_view(lse_tensor)
                     lse_row = lse_arr[cutlass.Int32(0), head_idx, :]
                     lse_row[_cu_q_b + q_row_global] = lse_val
             else:
-                if q_row_global < seqlen_q:
+                _row_valid = q_row_global < seqlen_q
+                if _row_valid:
                     lse_arr = cutlass.make_array_view(lse_tensor)
                     lse_row = lse_arr[batch_idx, head_idx, :]
                     lse_row[q_row_global] = lse_val
 
             sO_sub_base = sO[qs].base
 
+            _amax_o_ptr = Pointer(amax_o_tensor.iterator.raw_ptr(), dtype=cutlass.Int32)
+            _amax_o_local = cutlass.Float32(0.0)
+
             for chunk_idx in cutlass.range_constexpr(N_CHUNKS_O):
                 o_addr = tmem_base_epi + cutlass.Int32(tmem_O_off + chunk_idx * O_CHUNK)
                 o_chunk = nvvm.tcgen05_ld(
@@ -2147,6 +2154,9 @@ def _correction_warp_group(
                 )
                 nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD)
                 o_scaled = o_chunk * inv_sum
+                for _i in cutlass.range_constexpr(O_CHUNK):
+                    _e = o_scaled[_i]
+                    _amax_o_local = cute.math.max(_amax_o_local, cute.math.max(_e, -_e))
                 o_out = o_scaled.to(OUT_STORAGE_DTYPE)
 
                 col_offset_const = (chunk_idx * O_CHUNK) % D_BLOCK_SIZE
@@ -2160,6 +2170,10 @@ def _correction_warp_group(
                     bars.mb_o_empty[qs].wait(o_empty_phase)
                 smem_ptr.store_swizzled(o_out, alignment=64, swizzle=_O_SMEM_SWIZZLE)
 
+            # One atomic per valid row (invalid/OOB rows must not poison the global amax).
+            if _row_valid:
+                nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_o_ptr, _amax_o_local.bitcast(cutlass.Int32))
+
             # fence_proxy needed before TMA reads SMEM written by tcgen05_st.
             nvvm.fence_proxy("async.shared", space="cta")
 
@@ -2208,6 +2222,7 @@ def _host(
     sf_k_tensor: cute.Tensor,
     sf_v_tensor: cute.Tensor,
     lse_tensor: cute.Tensor,
+    amax_o_tensor: cute.Tensor,
     sinks_tensor: cute.Tensor,
     seq_kv_lens_tensor: cute.Tensor,
     o_desc_words: cute.Tensor,
@@ -2339,6 +2354,7 @@ def _build_sf_desc(sf_tensor, num_tiles, sf_smem_size, num_rows_box, num_heads):
         tma_k_sf_desc,
         tma_v_sf_desc,
         lse_tensor,
+        amax_o_tensor,
         sinks_tensor,
         seq_kv_lens_tensor,
         o_desc_words,
@@ -2432,6 +2448,12 @@ def compile(
         stride_order=(2, 1, 0),
         assumed_align=16,
     )
+    fake_amax_o = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float32,
+        (1,),
+        stride_order=(0,),
+        assumed_align=16,
+    )
     fake_sinks = cute.runtime.make_fake_compact_tensor(
         cutlass.Float32,
         (qh,),
@@ -2466,6 +2488,7 @@ def compile(
         fake_sf_k,
         fake_sf_v,
         fake_lse,
+        fake_amax_o,
         fake_sinks,
         fake_seq_kv_lens,
         fake_o_desc,
@@ -2475,5 +2498,5 @@ def compile(
         cutlass.Int32(_q_sf_tiles),
         cutlass.Int32(_kv_sf_tiles),
         stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False),
-        options="--enable-tvm-ffi",  # SM100: cfence not honored / not emitted
+        options="--enable-tvm-ffi",
     )