sdpa: fix SM100 frost zero-KV cluster deadlock + enable zero-length seq_len_kv sweeps

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

@@ -45,6 +45,10 @@ class Bars(NamedTuple):
     mb_empty_mainloop: object
 
     mb_q_o_alias: object
+    # Return edge of the Q∪O alias gate (see the soundness note in
+    # make_classic_bars): TMA-LDG arrives after consuming each alias-gate
+    # phase; TMA-STG waits it before its next alias arrive.
+    mb_qo_slab_free: object
 
 
 class D256Bars(NamedTuple):
@@ -140,7 +144,19 @@ def _alloc(n):
         mb_o_empty=MBarrier(_alloc(CFG.TILES_Q), stages=CFG.TILES_Q, init_count=CFG.ONE_WARP, producer=Producer.THREAD),
         mb_tmem_dealloc=MBarrier(_alloc(1), stages=1, init_count=CORR_LANES_TOTAL, producer=Producer.THREAD),
         mb_empty_mainloop=MBarrier(_alloc(1), stages=1, init_count=CORR_LANES_TOTAL, producer=Producer.LEADER, scope=Scope.LEADER),
+        # Q∪O alias gate FULL/EMPTY pair.  mb_q_o_alias alone is UNSOUND:
+        # mbarrier parity waits deadlock once a producer runs >= 2 phases
+        # ahead, and on EMPTY tiles (zero-KV varlen sequences) the
+        # corr -> STG -> alias-arrive chain has NO dependency on TMA-LDG, so
+        # a delayed LDG warp loses the race and its bootstrap parity credit
+        # is consumed by a real arrive (observed: LDG parked forever at the
+        # tile-1 alias wait with the barrier already in phase 1, deadlocking
+        # the whole cluster).  mb_qo_slab_free is the return edge: LDG
+        # arrives it right after consuming each alias phase and STG waits it
+        # before each alias arrive, bounding either side's lead to one phase
+        # by construction.
         mb_q_o_alias=MBarrier(_alloc(CFG.TILES_Q), stages=CFG.TILES_Q, init_count=CFG.ONE_WARP, producer=Producer.THREAD),
+        mb_qo_slab_free=MBarrier(_alloc(CFG.TILES_Q), stages=CFG.TILES_Q, init_count=CFG.ONE_WARP, producer=Producer.THREAD),
     )
 
 

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

@@ -428,6 +428,7 @@ def _kernel(
                 bars.mb_o_empty[qs].init()
                 if cutlass.const_expr(IS_QO_ALIAS):
                     bars.mb_q_o_alias[qs].init()
+                    bars.mb_qo_slab_free[qs].init()
                 for c in cutlass.range_constexpr(CFG.N_BMM2_CHUNKS):
                     bars.mb_bmm2_ready[qs * CFG.N_BMM2_CHUNKS + c].init()
             for ks in cutlass.range_constexpr(CFG.STAGES_KV):
@@ -700,15 +701,24 @@ def _tmaldg_warp_group(
                 # TMA-STG advances the Q/O alias gate for every tile, including
                 # empty ones. Consume that transaction even though no Q reload
                 # is needed, or the next nonempty tile observes stale parity.
+                # mb_qo_slab_free is the RETURN edge: without it, runs of
+                # empty tiles let STG (whose empty-tile O store depends only
+                # on correction, never on this warp) race >= 2 alias phases
+                # ahead of a delayed LDG, and mbarrier parity waits deadlock
+                # at lead 2 — the observed zero-KV cluster hang.
                 _wait_mbarrier(mb_q_reload[0], q_empty_phase)
+                bars.mb_qo_slab_free[0].arrive()
                 _wait_mbarrier(mb_q_reload[1], q_empty_phase)
+                bars.mb_qo_slab_free[1].arrive()
                 q_empty_phase = q_empty_phase ^ 1
         else:
             # Prologue interleave Q[0] -> K[first] -> Q[1] -> V[first] -> mainloop —
             # K load starts before Q[1] is issued and V before kv mainloop.
             kv_row_base = kv_left * CFG.TILE_N
 
             _wait_mbarrier(mb_q_reload[0], q_empty_phase)
+            if cutlass.const_expr(IS_QO_ALIAS):
+                bars.mb_qo_slab_free[0].arrive()
             if cutlass.const_expr(CFG.CTA_MMA == 2):
                 bars.mb_q_full[0].arrive(n_bytes=qTmaTransactionBytes, pred=is_leader & nvvm.elect_sync())
             else:
@@ -744,6 +754,8 @@ def _tmaldg_warp_group(
             )
 
             _wait_mbarrier(mb_q_reload[1], q_empty_phase)
+            if cutlass.const_expr(IS_QO_ALIAS):
+                bars.mb_qo_slab_free[1].arrive()
             if cutlass.const_expr(CFG.CTA_MMA == 2):
                 bars.mb_q_full[1].arrive(n_bytes=qTmaTransactionBytes, pred=is_leader & nvvm.elect_sync())
             else:
@@ -862,6 +874,9 @@ def _tmastg_warp_group(
     via scheduler warp's clusterlaunchcontrol.try_cancel.async.
     """
     o_full_phase = cutlass.Int32(0)  # consumer waits — first-arrive flips 0 → 1
+    # Alias-gate return edge (see the arrive site below): starts 0 — the
+    # first wait consumes LDG's first real slab_free arrive.
+    slab_free_phase = cutlass.Int32(0)
 
     tma_o = GmemTileTma(tma_o_desc)
 
@@ -906,11 +921,22 @@ def _tmastg_warp_group(
 
             bars.mb_o_empty[qs].arrive()
             # QO_ALIAS: O[qs] has drained to GMEM → the shared Q∪O slab is free
-            # for TMA-LDG to clobber with the next tile's Q[qs].
+            # for TMA-LDG to clobber with the next tile's Q[qs].  The
+            # mb_qo_slab_free wait (return edge) throttles this arrive to at
+            # most ONE phase ahead of LDG's alias-gate consumption — without
+            # it, empty-tile runs (zero-KV varlen) let this warp lap a
+            # delayed LDG by 2+ phases and mbarrier parity waits deadlock.
+            # Consumer-side bootstrap: slab_free_phase starts 0, so the first
+            # wait consumes LDG's first REAL arrive (LDG cannot be preceded —
+            # its arrive follows its own alias wait, which this warp's tile-1
+            # arrive has not yet advanced).
             if cutlass.const_expr(IS_QO_ALIAS):
+                _wait_mbarrier(bars.mb_qo_slab_free[qs], slab_free_phase)
                 bars.mb_q_o_alias[qs].arrive()
 
         o_full_phase = o_full_phase ^ 1
+        if cutlass.const_expr(IS_QO_ALIAS):
+            slab_free_phase = slab_free_phase ^ 1
 
         _wait_ptr(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase)
         nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()

test/python/sdpa/random_config.py

@@ -355,11 +355,7 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None):
             # ~10% chance of 0-length sequence for each batch
             randoms_.seq_len_q = [0 if rng.random() < 0.1 else rng.randint(1, randoms_.s_q) for _ in range(randoms_.batches)]
             # ~10% chance of 0-length sequence for each batch (independent of seq_len_q)
-            randoms_.seq_len_kv = [
-                # 0 if rng.random() < 0.1 else rng.randint(randoms_.seq_len_q[i], randoms_.s_kv) for i in range(randoms_.batches)
-                rng.randint(1, randoms_.s_kv)
-                for i in range(randoms_.batches)
-            ]
+            randoms_.seq_len_kv = [0 if rng.random() < 0.1 else rng.randint(1, randoms_.s_kv) for i in range(randoms_.batches)]
 
         # Decide the left and right bounds for the sliding window mask (None = no bound)
         randoms_.left_bound = None