perf(rocke): gfx950 d128 prefill softmax-MFMA interleave + nw=4 dispatch

dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py

@@ -34,6 +34,7 @@
     _enable_mfma_32x32,
     _enable_register_pv,
     _enable_sched_barrier,
+    _enable_softmax_mfma_interleave,
     _enable_transposed_half_local_pv,
     _enable_transposed_qk_32x32,
     _enable_transposed_subflags,
@@ -205,6 +206,16 @@ def _tiled_spec_from_problem(
         )
     if "use_sched_barrier" in _spec_field_names:
         _gfx950_schedule_fields["use_sched_barrier"] = _enable_sched_barrier(problem)
+    # gfx950 d128 softmax<->MFMA interleave lever (iglp_opt(1)); paired with the
+    # nw=4 widening in _select_2d_num_warps for the same cohort. Field-presence
+    # guarded (gfx942/gfx1250 spec classes lack it). Mutually exclusive with
+    # use_sched_barrier -- the cohorts do not overlap (sched_barrier is the
+    # nw==1 short-prefill cohort; interleave is the wider d128 combo).
+    if "use_softmax_mfma_interleave" in _spec_field_names and (
+        _enable_softmax_mfma_interleave(problem)
+    ):
+        _gfx950_schedule_fields["use_softmax_mfma_interleave"] = True
+        _gfx950_schedule_fields["softmax_interleave_mode"] = 1
     # d128 long-context lever: K single-buffer lets the larger T=64 tile fit
     # the 2-WG/CU LDS budget at HD=128 (see _select_2d_tile_size). Gated on the
     # same d128 small-tile cohort + opt-in env so default/production routing is

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py

@@ -585,6 +585,81 @@ def _enable_k_single_buffer(problem: UnifiedAttentionProblem) -> bool:
     return _enable_d128_small_tile(problem) and problem.block_size >= 32
 
 
+def _enable_softmax_mfma_interleave(problem: UnifiedAttentionProblem) -> bool:
+    """gfx950 single-batch d128 prefill: interleave the softmax VALU into the
+    MFMA window via ``iglp_opt(1)`` and widen to num_warps=4.
+
+    The shipped selector routes this cohort to num_warps=2 with no MFMA/softmax
+    interleave, which leaves the Matrix cores starved: rocprofv3 on D128 S8192
+    (Llama-3.1-405B, bf16) shows MFMA-busy ~5% / occupancy ~12% at 2 waves. The
+    ``use_softmax_mfma_interleave`` lever (``iglp_opt(1)`` at the loop top, which
+    lets the backend spread the causal-mask + softmax VALU across the MFMA gap)
+    combined with num_warps=4 lifts MFMA-busy to ~18% and occupancy to ~24%.
+
+    Measured (gfx950 MI355X, comgr llvm22, same-session vs the shipped auto path;
+    ``sweep_shapes_gfx950.py`` across the gap-doc cohort, correctness max_abs=0.0
+    vs the shipped kernel on every shape):
+      * D128 S8192: 122 -> 443 TF (3.6x); S4096 118 -> 412 (3.5x);
+        S2048 109 -> 350 (3.2x); S1024 108 -> 278 (2.6x); S512 105 -> 198 (1.9x).
+      * Holds across Llama-3-8B/70B, Llama-3.1-405B, Qwen2.5-7B, Qwen3-235B,
+        GPT-3-13B (MHA); bf16 and fp16 alike (no dtype cliff).
+
+    Occupancy (probe_occupancy / llvm-readelf on the produced HSACO): the widened
+    nw=4 + interleave kernel is 250 VGPR / 0 AGPR / 0 spill / 64 KB LDS -> 8
+    waves/CU, strictly BETTER than the shipped nw=2 kernel (296 VGPR / 40 AGPR,
+    4 waves/CU) -- so this does not trip the ``_enable_d128_small_tile`` nw=2
+    occupancy reconciliation (that concern was T=32 on llvm20; here T stays 128).
+
+    Scoped to the exact single-batch d128 combo cohort. The residual gap to
+    FlyDSL (~1090 TF) is structural (an 8-wave/32x32 dual-wave rewrite) and is
+    tracked separately; this is the hint-only ceiling and ships standalone.
+
+    **Why this cohort only (measured, not assumed).** The interleave lever pays
+    off wherever the loop is schedule-bound with a softmax-VALU/MFMA-idle window
+    to fill -- i.e. a 32x32-transposed (combo) geometry at nw>=2. The full
+    eligible space was swept (MI355X); adding ONLY the hint to each cohort's
+    SHIPPED spec:
+
+        cohort (shipped geometry)        S1024  S2048  S4096
+        d128 single-batch combo, nw4      1.10x  1.05x  1.03x  <- the gated win
+        d64  single-batch combo, nw4      0.97x  0.95x  0.93x  (regresses)
+        d128 multi-seq combo, nw2         1.00x  1.00x  0.99x  (neutral)
+        d256 single-seq plain, nw1        1.00x  1.00x  0.99x  (see note)
+        d128 sliding-window plain, nw2    0.83x  0.72x  0.73x  (see note)
+
+    The hint is inert on d128 multi-seq (already nw2-combo, no idle window to
+    exploit at that occupancy) and regresses d64 and d128-sliding-window on
+    their shipped geometry. It matches the runbook (arch/gfx950.md): iglp_opt
+    is neutral-to-negative without a real MFMA-idle window.
+
+    **d256 and d128-sliding-window ALSO benefit from this lever -- covered
+    elsewhere, not dropped.** On their SHIPPED plain geometry the hint alone is
+    inert/negative, but with the combo+nw geometry those cohorts want, the lever
+    helps there too and is being enabled by the cohort-owning PRs: d256 gfx950
+    via PR #9233 (AICK-1495, 32x32 combo + interleave mode2/g4, 1.1-1.94x vs
+    AOTriton), sliding-window d128 via AICK-1492 (wide-atom windowed path, in
+    progress). This gate stays scoped to the single-batch d128 combo so it does
+    not collide with those; the lever's d256/SWA coverage lives with the kernel
+    geometry work that unlocks it.
+
+    ESCAPE HATCH: ``HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE=0`` (or off/no/false)
+    force-DISABLES the lever, restoring the prior nw=2 / no-interleave routing.
+    """
+    env = (
+        __import__("os")
+        .environ.get("HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE", "")
+        .strip()
+        .lower()
+    )
+    if env in ("0", "false", "no", "off"):
+        return False
+    return (
+        problem.head_size == 128
+        and not problem.use_fp8
+        and _enable_single_batch_combo(problem)
+    )
+
+
 def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int:
     """Choose ``tile_size`` (T) for the tiled 2D kernel.
 
@@ -824,6 +899,13 @@ def _select_2d_num_warps(problem: UnifiedAttentionProblem) -> int:
         # gate is OFF this branch is byte-for-byte the prior behavior.
         if problem.head_size == 64:
             target = 4
+        elif _enable_softmax_mfma_interleave(problem):
+            # d128 softmax-MFMA-interleave cohort: widen to nw=4. The interleaved
+            # kernel is occupancy-BETTER than the shipped nw=2 (250 vs 296 VGPR,
+            # 8 vs 4 waves/CU -- see _enable_softmax_mfma_interleave), so the
+            # small-tile nw=2 reconciliation below does not apply. Overrides the
+            # _enable_d128_small_tile -> nw=2 rule for this cohort only.
+            target = 4
         elif _enable_d128_small_tile(problem):
             target = 2
         elif problem.max_seqlen_q <= 1024:

dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c

@@ -829,6 +829,37 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s)
         s->use_transposed_half_local_pv = true;
         break;
 
+    /* idx53: production nw=4 softmax<->MFMA interleave config (idx52 +
+       use_softmax_mfma_interleave mode 1 + num_warps=4). Emits one iglp_opt(1)
+       after the iter-start K drain; ported in kv_body_qk_softmax.cpp, so it MUST
+       be C/Python byte-identical. Kept in lockstep with the Python _CONFIGS. */
+    case 53:
+        s->head_size = 128;
+        s->block_size = 64;
+        s->num_query_heads = 64;
+        s->num_kv_heads = 8;
+        s->dtype = "bf16";
+        s->use_sinks = false;
+        s->sliding_window = 0;
+        s->has_softcap = false;
+        s->num_seqs = 1;
+        s->num_warps = 4;
+        s->block_m_per_warp = 32;
+        s->has_tile_size = true;
+        s->tile_size = 128;
+        s->use_k_single_buffer = true;
+        s->use_mfma_32x32 = true;
+        s->use_transposed_qk_32x32 = true;
+        s->use_transposed_scalar_state = true;
+        s->use_transposed_invariant_hoist = true;
+        s->use_transposed_mask_once = true;
+        s->use_transposed_mask_limit = true;
+        s->use_mfma32_skip_legacy_qreg = true;
+        s->use_transposed_half_local_pv = true;
+        s->use_softmax_mfma_interleave = true;
+        s->softmax_interleave_mode = 1;
+        break;
+
     default:
         return -1;
     }

dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py

@@ -816,6 +816,37 @@
         use_mfma32_skip_legacy_qreg=True,
         use_transposed_half_local_pv=True,
     ),
+    # idx53: the production nw=4 softmax<->MFMA interleave config. This is exactly
+    # what the gfx950 single-batch d128 selector now routes to (see
+    # _enable_softmax_mfma_interleave): block_size=64 -> tile_size=128 (2*BS) so
+    # BLOCK_M=128 (nw4 * mw32) fits the single K slot (block_m <= tile_size). The
+    # interleave emits one iglp_opt(1) after the iter-start K drain; the gfx950 C
+    # twin ports it (kv_body_qk_softmax.cpp), so it MUST be C/Python byte-identical.
+    53: dict(
+        head_size=128,
+        block_size=64,
+        num_query_heads=64,
+        num_kv_heads=8,
+        dtype="bf16",
+        use_sinks=False,
+        sliding_window=0,
+        has_softcap=False,
+        num_seqs=1,
+        num_warps=4,
+        block_m_per_warp=32,
+        tile_size=128,
+        use_k_single_buffer=True,
+        use_mfma_32x32=True,
+        use_transposed_qk_32x32=True,
+        use_transposed_scalar_state=True,
+        use_transposed_invariant_hoist=True,
+        use_transposed_mask_once=True,
+        use_transposed_mask_limit=True,
+        use_mfma32_skip_legacy_qreg=True,
+        use_transposed_half_local_pv=True,
+        use_softmax_mfma_interleave=True,
+        softmax_interleave_mode=1,
+    ),
 }
 
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h

@@ -181,6 +181,10 @@ bool rocke_unified_attn_enable_transposed_subflags(const rocke_unified_attn_prob
  * (T=block_size + nw=2 -> 2 WG/CU). DEFAULT-ON for the gfx950 single-batch
  * d128 no-FP8 combo; HIPDNN_GFX950_D128_SMALL_TILE=0 force-disables. */
 bool rocke_unified_attn_enable_d128_small_tile(const rocke_unified_attn_problem_t* p);
+/* Python: _enable_softmax_mfma_interleave(problem) -- d128 softmax<->MFMA
+ * interleave (iglp_opt(1)) + nw=4. DEFAULT-ON for the gfx950 single-batch d128
+ * no-FP8 combo; HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE=0 force-disables. */
+bool rocke_unified_attn_enable_softmax_mfma_interleave(const rocke_unified_attn_problem_t* p);
 /* Python: _enable_v_double_buffer(problem) -- short single-batch combo prefill. */
 bool rocke_unified_attn_enable_v_double_buffer(const rocke_unified_attn_problem_t* p);
 /* Python: _enable_early_v_schedule(problem) -- long single-batch combo prefill. */

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h

@@ -133,6 +133,10 @@ typedef struct rocke_gfx950_attn2d_build_ctx
     /* Lever-3 (CK-Tile-derived) sched_barrier fence after the QK MFMA cluster. */
     bool USE_SCHED_BARRIER;
     int SCHED_BARRIER_MASK;
+    /* STEP 2 lever: softmax<->MFMA interleave via iglp_opt at the loop top.
+     * modes 0/1 emit iglp_opt(mode); mutually exclusive with USE_SCHED_BARRIER. */
+    bool USE_SOFTMAX_INTERLEAVE;
+    int SOFTMAX_INTERLEAVE_MODE;
     /* fp8 K/V cache predicates */
     bool KV_FP8, FP8_MFMA_QK, FP8_MFMA_PV;
     bool REGISTER_PV, USE_MFMA_32X32, TRANSPOSED_QK_32X32;

dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp

@@ -52,6 +52,48 @@ static bool arch_is(const char* want)
     return strcmp(rocke_unified_attn_resolve_arch(), want) == 0;
 }
 
+/* ----------------------------------------------------- env off-switch gate */
+
+/* Mirror Python's ``os.environ.get(name, "").strip().lower() in
+ * ("0","false","no","off")`` for a boolean feature kill-switch. Returns true
+ * when the env var is set to a recognized OFF keyword (feature force-disabled),
+ * false otherwise (unset, empty, or any other value -> feature stays enabled).
+ *
+ * Strips leading/trailing ASCII whitespace and case-folds before comparison so
+ * values like ``"0 "`` / ``" off"`` behave identically to the Python gate. */
+static bool rocke_env_forces_off(const char* name)
+{
+    const char* env = getenv(name);
+    if(env == NULL)
+    {
+        return false;
+    }
+    /* skip leading whitespace (Python str.strip default set) */
+    while(*env == ' ' || *env == '\t' || *env == '\n' || *env == '\r' || *env == '\v'
+          || *env == '\f')
+    {
+        ++env;
+    }
+    /* copy + case-fold up to the first trailing-whitespace / capacity bound */
+    char buf[16];
+    size_t n = 0;
+    for(; env[n] != '\0' && n + 1 < sizeof(buf); ++n)
+    {
+        char c = env[n];
+        buf[n] = (c >= 'A' && c <= 'Z') ? (char)(c - 'A' + 'a') : c;
+    }
+    buf[n] = '\0';
+    /* trim trailing whitespace */
+    while(n > 0
+          && (buf[n - 1] == ' ' || buf[n - 1] == '\t' || buf[n - 1] == '\n' || buf[n - 1] == '\r'
+              || buf[n - 1] == '\v' || buf[n - 1] == '\f'))
+    {
+        buf[--n] = '\0';
+    }
+    return strcmp(buf, "0") == 0 || strcmp(buf, "false") == 0 || strcmp(buf, "no") == 0
+           || strcmp(buf, "off") == 0;
+}
+
 /* ----------------------------------------------------- num_queries_per_kv */
 
 int rocke_unified_attn_num_queries_per_kv(rocke_ir_builder_t* b,
@@ -177,24 +219,26 @@ static bool enable_single_batch_combo(const rocke_unified_attn_problem_t* p)
  * the Python helper's env gate byte-faithfully (case-folded, off-keywords only). */
 static bool enable_d128_small_tile(const rocke_unified_attn_problem_t* p)
 {
-    const char* env = getenv("HIPDNN_GFX950_D128_SMALL_TILE");
-    if(env != NULL)
-    {
-        /* case-fold then compare to the OFF keywords (matches Python
-           .strip().lower() in ("0","false","no","off")). */
-        char buf[16];
-        size_t i = 0;
-        for(; env[i] != '\0' && i + 1 < sizeof(buf); ++i)
-        {
-            char c = env[i];
-            buf[i] = (c >= 'A' && c <= 'Z') ? (char)(c - 'A' + 'a') : c;
-        }
-        buf[i] = '\0';
-        if(strcmp(buf, "0") == 0 || strcmp(buf, "false") == 0 || strcmp(buf, "no") == 0
-           || strcmp(buf, "off") == 0)
-        {
-            return false;
-        }
+    /* ESCAPE HATCH: HIPDNN_GFX950_D128_SMALL_TILE=0 (or off/no/false)
+       force-DISABLES it. Mirrors the Python .strip().lower() env gate. */
+    if(rocke_env_forces_off("HIPDNN_GFX950_D128_SMALL_TILE"))
+    {
+        return false;
+    }
+    return p->head_size == 128 && !p->use_fp8 && enable_single_batch_combo(p);
+}
+
+/* Python: _enable_softmax_mfma_interleave(problem). gfx950 single-batch d128
+ * prefill: enable use_softmax_mfma_interleave (iglp_opt(1)) and widen to nw=4.
+ * The widened kernel is occupancy-BETTER than the shipped nw=2 (250 vs 296 VGPR,
+ * 8 vs 4 waves/CU), so it does not trip the d128 small-tile nw=2 reconciliation.
+ * ESCAPE HATCH: HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE=0 (or off/no/false)
+ * force-DISABLES it. Mirrors the Python env gate byte-faithfully. */
+static bool enable_softmax_mfma_interleave(const rocke_unified_attn_problem_t* p)
+{
+    if(rocke_env_forces_off("HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE"))
+    {
+        return false;
     }
     return p->head_size == 128 && !p->use_fp8 && enable_single_batch_combo(p);
 }
@@ -390,6 +434,16 @@ static int select_2d_tile_size(const rocke_unified_attn_problem_t* p)
         {
             return 128;
         }
+        /* d128 softmax-MFMA-interleave cohort keeps the WIDE T=2*BS tile (the
+           nw=4 interleave win is measured at the wide tile; nw=4 + T=32 is the
+           occupancy-WORSE combo -- 56 KB LDS -> 1 WG/CU). Guard BEFORE the
+           small-tile check so this cohort mirrors Python _select_2d_tile_size
+           (which returns 2*block_size here). Matches the num_warps selector's
+           interleave-before-small-tile ordering. */
+        if(enable_softmax_mfma_interleave(p))
+        {
+            return 2 * p->block_size;
+        }
         /* d128 occupancy lever: halve the tile to T=block_size (one paged
            block per iter) -> LDS 48->24 KB -> 2 WG/CU. DEFAULT-ON;
            HIPDNN_GFX950_D128_SMALL_TILE=0 force-disables. */
@@ -477,6 +531,15 @@ int rocke_unified_attn_select_2d_num_warps(const rocke_unified_attn_problem_t* p
         {
             t2 = 4;
         }
+        /* d128 softmax-MFMA-interleave cohort: widen to nw=4. The interleaved
+           kernel is occupancy-BETTER than the shipped nw=2 (250 vs 296 VGPR,
+           8 vs 4 waves/CU -- see enable_softmax_mfma_interleave), so the
+           small-tile nw=2 reconciliation below does not apply. Overrides the
+           _enable_d128_small_tile -> nw=2 rule for this cohort only. */
+        else if(enable_softmax_mfma_interleave(p))
+        {
+            t2 = 4;
+        }
         /* RECONCILIATION: the d128 small-tile occupancy win (T=block_size
            -> 2 WG/CU) REQUIRES num_warps=2 for ALL seqlens; nw=4 + T=32 is
            occupancy-WORSE (56 KB LDS -> 1 WG/CU). Override the S>=2048 -> nw=4
@@ -744,6 +807,11 @@ bool rocke_unified_attn_enable_d128_small_tile(const rocke_unified_attn_problem_
     return enable_d128_small_tile(p);
 }
 
+bool rocke_unified_attn_enable_softmax_mfma_interleave(const rocke_unified_attn_problem_t* p)
+{
+    return enable_softmax_mfma_interleave(p);
+}
+
 bool rocke_unified_attn_enable_v_double_buffer(const rocke_unified_attn_problem_t* p)
 {
     return enable_v_double_buffer(p);

dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp

@@ -587,6 +587,14 @@ void rocke_gfx950_attn2d_emit_kv_body(rocke_gfx950_attn2d_build_ctx_t* ctx)
     rocke_b_s_waitcnt(b, /*vmcnt=*/0, /*lgkmcnt=*/0, /*expcnt=*/-1);
     rocke_b_sync(b);
 
+    /* STEP 2 lever (modes 0/1): canned iglp interleave for the whole loop body,
+     * placed once after the iter-start K drain (mirrors Python 3093-3094). Level
+     * 0 = GEMM MFMA-interleave, level 1 = attention-style interleave. Mode 2
+     * (sched_group_barrier grouping) is rejected in the glue. */
+    if(ctx->USE_SOFTMAX_INTERLEAVE
+       && (ctx->SOFTMAX_INTERLEAVE_MODE == 0 || ctx->SOFTMAX_INTERLEAVE_MODE == 1))
+        rocke_b_iglp_opt(b, ctx->SOFTMAX_INTERLEAVE_MODE);
+
     /* Early-V schedule: V_lds is single-buffered and the iter-start full drain
      * guarantees the previous PV's V reads retired, so the current V copy can be
      * issued before QK to overlap with QK + softmax (Python 2566-2570). */