feat(hip-kernel-provider): port the rocKE gfx942 tiled-2D mask-limit, sliced-K ring, phase split and HBM-direct K/V to C++

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

@@ -26,9 +26,6 @@
 
 # Python spec fields the shared C tiled-2D struct has no slot for.
 _UNPORTED_2D = {
-    "use_v_hbm_direct",
-    "use_k_hbm_direct",
-    "use_causal_mask_phase_split",
     "use_additive_bias",
     "causal_top_left",
     "right_bound",
@@ -72,6 +69,24 @@ def test_gfx942_2d_binding_keeps_the_x8_path(engine):
     assert actual == expected
 
 
+@pytest.mark.parametrize(
+    "field", ["use_v_hbm_direct", "use_k_hbm_direct", "use_causal_mask_phase_split"]
+)
+def test_gfx942_2d_binding_reads_gfx942_only_fields(engine, field):
+    spec = Gfx942Spec2D(
+        **_REQUIRED_2D,
+        dtype="fp16",
+        block_m_per_warp=32,
+        tile_size=64,
+        use_mfma_32x32x8=True,
+        use_transposed_qk_32x32=True,
+        **{field: True},
+    )
+    expected = _lower_python(build_gfx942_2d(spec, arch="gfx942"), "gfx942")
+    actual = engine.gfx942_attention_tiled_2d_lower_llvm(asdict(spec), "gfx942")
+    assert actual == expected
+
+
 def test_gfx950_2d_default_spec_round_trips(engine):
     spec = Gfx950Spec2D(**_REQUIRED_2D, dtype="bf16")
     expected = _lower_python(build_gfx950_2d(spec, arch="gfx950"), "gfx950")

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

@@ -21,7 +21,7 @@
     UnifiedAttentionProblem,
     _canonical_arch_for_selection,
     _enable_gfx942_flash_k_sliced_ring,
-    _enable_gfx942_flash_mask_limit,
+    _enable_i64_kv_addr,
     _gfx942_4warp_fast,
     _mask_label,
     _reject_mask_on_arch,
@@ -253,10 +253,11 @@ def _tiled_2d_not_ported(p: UnifiedAttentionProblem, arch: str):
     if arch == "gfx942":
         if _gfx942_4warp_fast(p, arch):
             return "gfx942 4-warp GQA tiled 2D is not ported to C++ yet"
-        if _enable_gfx942_flash_k_sliced_ring(p, arch):
-            return "gfx942 tiled 2D sliced-K ring is not ported to C++ yet"
-        if _enable_gfx942_flash_mask_limit(p, arch):
-            return "gfx942 tiled 2D flash mask-limit is not ported to C++ yet"
+        if _enable_i64_kv_addr(p) and _enable_gfx942_flash_k_sliced_ring(p, arch):
+            return (
+                "gfx942 tiled 2D i64 KV addressing with the sliced-K ring is not "
+                "ported to C++ yet"
+            )
         return None
     if p.use_additive_bias:
         return "tiled 2D additive bias is not ported to C++ yet"

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

@@ -237,13 +237,16 @@ def test_selectors_follow_env(engine, monkeypatch, env):
         "use_conflict_free_v_store",
         "use_conflict_free_v_store_split",
         "use_global_load_lds_k",
+        "use_causal_mask_phase_split",
         "use_iglp_opt",
+        "use_k_hbm_direct",
         "use_k_sliced_ldsseq",
         "use_k_sliced_ring",
         "use_mfma_32x32x8",
         "use_q_direct_global",
         "use_q_major_grid",
         "use_qk_pv_sched_group_barrier",
+        "use_v_hbm_direct",
     }
 )
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine_attention.cpp

@@ -844,9 +844,6 @@ std::vector<std::string> sage_verify(const py::dict& d, const std::string& arch)
 /* Python fields of the gfx942 and gfx950 UnifiedAttention2DTiledSpec that the
  * shared C spec struct has no slot for. */
 const unported_field k_t2d_unported[] = {
-    {"use_v_hbm_direct", false, 0},
-    {"use_k_hbm_direct", false, 0},
-    {"use_causal_mask_phase_split", false, 0},
     {"use_additive_bias", false, 0},
     {"causal_top_left", false, 0},
     {"right_bound", false, 0},
@@ -937,6 +934,10 @@ void t2d_read(const py::dict& d, Store& st, rocke_attention_tiled_2d_spec_t& s)
     s.use_kq_lds_pad = a_bool(d, "use_kq_lds_pad", s.use_kq_lds_pad);
     s.kq_lds_pad_halves = a_int(d, "kq_lds_pad_halves", s.kq_lds_pad_halves);
     s.use_mask_phase_split = a_bool(d, "use_mask_phase_split", s.use_mask_phase_split);
+    s.use_v_hbm_direct = a_bool(d, "use_v_hbm_direct", s.use_v_hbm_direct);
+    s.use_k_hbm_direct = a_bool(d, "use_k_hbm_direct", s.use_k_hbm_direct);
+    s.use_causal_mask_phase_split
+        = a_bool(d, "use_causal_mask_phase_split", s.use_causal_mask_phase_split);
     std::string v;
     if(a_str(d, "dtype", v))
         s.dtype = st.keep(v);
@@ -1567,6 +1568,9 @@ static py::dict spec_to_dict(const rocke_attention_tiled_2d_spec_t& s)
     d["use_kq_lds_pad"] = s.use_kq_lds_pad;
     d["kq_lds_pad_halves"] = s.kq_lds_pad_halves;
     d["use_mask_phase_split"] = s.use_mask_phase_split;
+    d["use_v_hbm_direct"] = s.use_v_hbm_direct;
+    d["use_k_hbm_direct"] = s.use_k_hbm_direct;
+    d["use_causal_mask_phase_split"] = s.use_causal_mask_phase_split;
     return d;
 }
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h

@@ -177,6 +177,13 @@ typedef struct rocke_attention_tiled_2d_spec
     /* gfx950-only: split the KV loop into a mask-free full-tile phase and a
      * masked boundary phase (transposed-32x32 mask-limit path, no SW). */
     bool use_mask_phase_split; /* False */
+    /* gfx942 only (the gfx950 builder ignores them, as it does the ring and
+     * cfvst fields): read V / K straight from HBM in the PV / QK MFMA loops
+     * instead of staging them through LDS, and split the causal KV loop into an
+     * unmasked full-tile phase plus a masked boundary phase. */
+    bool use_v_hbm_direct; /* False */
+    bool use_k_hbm_direct; /* False */
+    bool use_causal_mask_phase_split; /* False */
 } rocke_attention_tiled_2d_spec_t;
 
 /* Materialise every defaulted field (the dataclass defaults). The required

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

@@ -125,6 +125,7 @@ typedef struct rocke_gfx942_attn2d_build_ctx
     bool TRANSPOSED_SCALAR_STATE, TRANSPOSED_INVARIANT_HOIST;
     bool TRANSPOSED_MASK_ONCE, TRANSPOSED_HALF_LOCAL_PV;
     bool SKIP_LEGACY_QREG, TRANSPOSED_MASK_LIMIT, GROUPED_KV2;
+    bool CAUSAL_MASK_PHASE_SPLIT, V_HBM_DIRECT, K_HBM_DIRECT;
     bool FAST_PAGED_KV_DESC, I64_KV_ADDR, EARLY_V_SCHEDULE;
     /* fp8 K/V cache predicates */
     bool KV_FP8, FP8_MFMA_QK, FP8_MFMA_PV, FP8_NATIVE_QK;
@@ -332,6 +333,12 @@ typedef struct rocke_gfx942_attn2d_build_ctx
     rocke_value_t* hoist_q_pos[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE];
     rocke_value_t* hoist_q_head[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE];
     rocke_value_t* hoist_row_mask[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE];
+    /* Transposed-32x32 per-lane invariants hoisted before the KV loop under
+     * TRANSPOSED_INVARIANT_HOIST (NULL otherwise). */
+    rocke_value_t* st_qp_hoist;
+    rocke_value_t* st_row_ok_hoist;
+    rocke_value_t* st_causal_lim_hoist;
+    rocke_value_t* st_alibi_slope_hoist;
 
     /* ---- KV-loop body live carry (set per _emit_kv_body invocation) ---- *
      * The current loop iter var + the unpacked m/l/acc carry the inner QK/mask/
@@ -447,6 +454,12 @@ rocke_value_t* rocke_gfx942_attn2d_v_t_load(rocke_gfx942_attn2d_build_ctx_t* ctx
                                             rocke_value_t* dim,
                                             rocke_value_t* tok,
                                             int n);
+/* _v_hbm_elem: one V[tok, dim] element of KV tile `tile` read straight from
+ * HBM (V_HBM_DIRECT), zero past the valid KV length. */
+rocke_value_t* rocke_gfx942_attn2d_v_hbm_elem(rocke_gfx942_attn2d_build_ctx_t* ctx,
+                                              rocke_value_t* tile,
+                                              rocke_value_t* tok,
+                                              rocke_value_t* dim);
 rocke_value_t* rocke_gfx942_attn2d_v_load1(rocke_gfx942_attn2d_build_ctx_t* ctx,
                                            rocke_value_t* v_buf,
                                            rocke_value_t* v_row,

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

@@ -160,16 +160,13 @@ static bool reject_tiled_2d_not_ported(const rocke_unified_attention_problem_t*
             gate_reason(reason, reason_cap, "gfx942 4-warp GQA tiled 2D is not ported to C++ yet");
             return true;
         }
-        if(rocke_unified_attn_enable_gfx942_flash_k_sliced_ring(p, arch))
+        if(rocke_unified_attn_enable_i64_kv_addr(p)
+           && rocke_unified_attn_enable_gfx942_flash_k_sliced_ring(p, arch))
         {
-            gate_reason(
-                reason, reason_cap, "gfx942 tiled 2D sliced-K ring is not ported to C++ yet");
-            return true;
-        }
-        if(rocke_unified_attn_enable_gfx942_flash_mask_limit(p, arch))
-        {
-            gate_reason(
-                reason, reason_cap, "gfx942 tiled 2D flash mask-limit is not ported to C++ yet");
+            gate_reason(reason,
+                        reason_cap,
+                        "gfx942 tiled 2D i64 KV addressing with the sliced-K ring is not ported "
+                        "to C++ yet");
             return true;
         }
         return false;

dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d.cpp

@@ -131,6 +131,9 @@ rocke_attention_tiled_2d_spec_t rocke_attention_tiled_2d_spec_default(void)
     s.use_kq_lds_pad = false;
     s.kq_lds_pad_halves = 8;
     s.use_mask_phase_split = false;
+    s.use_v_hbm_direct = false;
+    s.use_k_hbm_direct = false;
+    s.use_causal_mask_phase_split = false;
     return s;
 }
 
@@ -536,6 +539,34 @@ bool rocke_attention_tiled_2d_spec_validate(rocke_ir_builder_t* b,
         return false;
     }
 
+    if(s->use_k_hbm_direct)
+    {
+        if(!(s->use_mfma_32x32x8 && s->use_transposed_qk_32x32))
+        {
+            rocke_attn2d_set_err(b,
+                                 ROCKE_ERR_VALUE,
+                                 "use_k_hbm_direct requires transposed-x8 "
+                                 "(use_mfma_32x32x8 + use_transposed_qk_32x32); "
+                                 "the 16x16 QK loop always reads K_lds");
+            return false;
+        }
+        if(s->use_k_sliced_ring)
+        {
+            rocke_attn2d_set_err(b,
+                                 ROCKE_ERR_VALUE,
+                                 "use_k_hbm_direct is mutually exclusive with use_k_sliced_ring");
+            return false;
+        }
+        if(s->use_global_load_lds_k)
+        {
+            rocke_attn2d_set_err(b,
+                                 ROCKE_ERR_VALUE,
+                                 "use_k_hbm_direct is mutually exclusive with "
+                                 "use_global_load_lds_k");
+            return false;
+        }
+    }
+
     if(s->use_global_load_lds_k && s->kv_storage_dtype != NULL)
     {
         rocke_attn2d_set_err(