ROCm · aosewski · Jun 3, 2026 · May 11, 2026 · May 11, 2026 · Jun 2, 2026
@@ -82,7 +82,7 @@ typename T::Kargs kargs{
     {},  // deterministic
     /* batch-mode: 8 stride fields */
 };
-if constexpr(K.has_mask) {
+if constexpr(hasMask(K)) {
     kargs.window_size_left  = -1;
     kargs.window_size_right = 0;  // causal
     kargs.mask_type = MASK_FROM_TOP_LEFT;
@@ -102,8 +102,9 @@ args.scalars[S::MASK_TYPE].i32         = static_cast<int>(
 ```
 
 The device bridge reads these scalars and assigns them to `kargs` after
-the aggregate init -- the spec-time `has_mask` flag only gates which
-inheritance base is active, not the runtime mask shape.
+the aggregate init -- the spec-time `mask_type` enum only gates which
+inheritance base is active (NO_MASK vs any other family), not the
+runtime mask shape itself.
 
 ### 1D Tensor Stride Convention
 

@@ -33,7 +33,6 @@
 #include <cstdio>
 #include <cstdlib>
 #include <numeric>
-#include <string_view>
 #include <vector>
 
 // Namespace aliases for named slot constants.
@@ -865,36 +864,26 @@ static bool runDqDkDvBatchVariant(const rocm_ck::FmhaBwdDQDKDVVariant& variant,
     assert(batch > 0);
     args.scalars[DKV::BATCH_SIZE].i32 = static_cast<int32_t>(batch);
 
-    // Mask geometry is variant-specific. The compiled spec is identical for
-    // _cmask, _cmask_br, and _swa — geometry is selected at runtime via the
-    // WINDOW_SIZE_LEFT/RIGHT and MASK_TYPE scalar slots.
-    // GenericAttentionMaskEnum: 0=NO_MASK, 1=MASK_FROM_TOP_LEFT,
-    //                           2=MASK_FROM_BOTTOM_RIGHT, 3=MASK_GENERIC.
-    // Use ends_with (not find) so a future variant whose name merely contains
-    // "_swa" or "_cmask_br" as a non-suffix substring won't be misclassified.
-    if(variant.spec.has_mask)
+    // Mask geometry is driven entirely by the spec's mask_type enum (AE-1).
+    // The compiled spec is shared across _cmask, _cmask_br, and _swa --
+    // family selection happens at runtime via MASK_TYPE, and the window
+    // pair distinguishes causal (unlimited left, no future tokens) from
+    // sliding-window attention.
+    if(rocm_ck::hasMask(variant.spec))
     {
-        const std::string_view name(variant.name);
-        if(name.ends_with("_swa"))
+        args.scalars[DKV::MASK_TYPE].i32 = static_cast<int32_t>(variant.spec.mask_type);
+        if(variant.spec.mask_type == rocm_ck::FmhaMaskType::GENERIC)
         {
-            // Sliding-window attention with both-side limits.
+            // Sliding-window attention. Window size is hardcoded today;
+            // surfacing it on the spec is tracked as a follow-up (AE-3).
             args.scalars[DKV::WINDOW_SIZE_LEFT].i32  = 64;
             args.scalars[DKV::WINDOW_SIZE_RIGHT].i32 = 64;
-            args.scalars[DKV::MASK_TYPE].i32         = 3; // MASK_GENERIC
-        }
-        else if(name.ends_with("_cmask_br"))
-        {
-            // Bottom-right causal: unlimited left, no future tokens.
-            args.scalars[DKV::WINDOW_SIZE_LEFT].i32  = -1;
-            args.scalars[DKV::WINDOW_SIZE_RIGHT].i32 = 0;
-            args.scalars[DKV::MASK_TYPE].i32         = 2; // MASK_FROM_BOTTOM_RIGHT
         }
         else
         {
-            // Default: top-left causal (preserves _cmask and _cmask_det).
+            // TOP_LEFT_CAUSAL and BOTTOM_RIGHT_CAUSAL share the causal window.
             args.scalars[DKV::WINDOW_SIZE_LEFT].i32  = -1;
             args.scalars[DKV::WINDOW_SIZE_RIGHT].i32 = 0;
-            args.scalars[DKV::MASK_TYPE].i32         = 1; // MASK_FROM_TOP_LEFT
         }
     }
 
@@ -1413,7 +1402,7 @@ int main(int argc, char** argv)
         // no deterministic flag get full verification.
         // Others are compilation proof only.
         bool is_plain_batch =
-            (!v.spec.has_mask && !v.spec.has_dropout && !v.spec.is_deterministic &&
+            (!rocm_ck::hasMask(v.spec) && !v.spec.has_dropout && !v.spec.is_deterministic &&
              v.spec.bias_type == rocm_ck::FmhaBiasType::NONE);
 
         bool passed = runDqDkDvBatchVariant(v,

@@ -187,21 +187,20 @@
         "hdim_q": 128,
         "hdim_v": 128,
         "mode": "batch",
-        "has_mask": True,
+        "mask_type": "top_left_causal",
         "block_size": 256,
     },
-    # _cmask_br and _swa share the compiled spec with _cmask. They are packed
-    # under distinct names so the loader can select mask geometry by variant
-    # name; the actual top-left/bottom-right/sliding-window selection happens
-    # at runtime via the WINDOW_SIZE_LEFT/RIGHT and MASK_TYPE scalar slots.
+    # _cmask_br and _swa share the compiled spec shape with _cmask but are
+    # disambiguated by mask_type. The window pair is set at runtime per
+    # family (causal: -1/0; sliding-window: 64/64).
     {
         "name": "fmha_bwd_dqdkdv_fp16_d128_batch_cmask_br",
         "family": "dqdkdv",
         "dtype": "fp16",
         "hdim_q": 128,
         "hdim_v": 128,
         "mode": "batch",
-        "has_mask": True,
+        "mask_type": "bottom_right_causal",
         "block_size": 256,
     },
     {
@@ -211,7 +210,7 @@
         "hdim_q": 128,
         "hdim_v": 128,
         "mode": "batch",
-        "has_mask": True,
+        "mask_type": "generic",
         "block_size": 256,
     },
     {
@@ -249,7 +248,7 @@
         "hdim_q": 128,
         "hdim_v": 128,
         "mode": "group",
-        "has_mask": True,
+        "mask_type": "top_left_causal",
         "block_size": 256,
     },
     {
@@ -382,7 +381,7 @@
         "hdim_q": 128,
         "hdim_v": 128,
         "mode": "batch",
-        "has_mask": True,
+        "mask_type": "top_left_causal",
         "is_deterministic": True,
         "block_size": 256,
     },

@@ -156,28 +156,24 @@ static constexpr FmhaBwdDQDKDVVariant ALL_DQDKDV_VARIANTS[] = {
          .signature = {.dtype = DataType::FP16,
                        .hdim_q = 128, .hdim_v = 128,
                        .mode = FmhaMode::BATCH},
-         .algorithm = {.has_mask = true,
+         .algorithm = {.mask_type = FmhaMaskType::TOP_LEFT_CAUSAL,
                        .pad_hdim_q = 8, .pad_hdim_v = 8}})},
-    // Bottom-right causal: same compiled spec as _cmask. The mask_type is
-    // selected at runtime via args.scalars[fmha_bwd_dqdkdv_slots::MASK_TYPE].
-    //
-    // Lookup note: findVariant() matches on spec features alone, so it returns
-    // _cmask first for any has_mask=true query. _cmask_br and _swa are
-    // reachable only via fmha_bwd_dqdkdv_variant_spec("<exact name>") (consteval)
-    // or by iterating ALL_DQDKDV_VARIANTS (host).
+    // Bottom-right causal: same compiled spec shape as _cmask, distinguished
+    // from it by mask_type so findVariant() returns each family unambiguously.
+    // The runtime MASK_TYPE scalar slot mirrors the spec field on launch.
     {"fmha_bwd_dqdkdv_fp16_d128_batch_cmask_br", makeSpec(FmhaBwdDQDKDVConfig{
          .signature = {.dtype = DataType::FP16,
                        .hdim_q = 128, .hdim_v = 128,
                        .mode = FmhaMode::BATCH},
-         .algorithm = {.has_mask = true,
+         .algorithm = {.mask_type = FmhaMaskType::BOTTOM_RIGHT_CAUSAL,
                        .pad_hdim_q = 8, .pad_hdim_v = 8}})},
-    // Sliding-window attention: same compiled spec as _cmask. window_size_left
-    // and window_size_right are runtime-parametrized via scalar slots.
+    // Sliding-window attention: window_size_left/right are runtime-parametrized
+    // via scalar slots; the spec only carries the mask family (GENERIC).
     {"fmha_bwd_dqdkdv_fp16_d128_batch_swa", makeSpec(FmhaBwdDQDKDVConfig{
          .signature = {.dtype = DataType::FP16,
                        .hdim_q = 128, .hdim_v = 128,
                        .mode = FmhaMode::BATCH},
-         .algorithm = {.has_mask = true,
+         .algorithm = {.mask_type = FmhaMaskType::GENERIC,
                        .pad_hdim_q = 8, .pad_hdim_v = 8}})},
     {"fmha_bwd_dqdkdv_fp16_d128_batch_det", makeSpec(FmhaBwdDQDKDVConfig{
          .signature = {.dtype = DataType::FP16,
@@ -199,7 +195,7 @@ static constexpr FmhaBwdDQDKDVVariant ALL_DQDKDV_VARIANTS[] = {
          .signature = {.dtype = DataType::FP16,
                        .hdim_q = 128, .hdim_v = 128,
                        .mode = FmhaMode::GROUP},
-         .algorithm = {.has_mask = true,
+         .algorithm = {.mask_type = FmhaMaskType::TOP_LEFT_CAUSAL,
                        .pad_hdim_q = 8, .pad_hdim_v = 8}})},
     {"fmha_bwd_dqdkdv_fp16_d128_group_det", makeSpec(FmhaBwdDQDKDVConfig{
          .signature = {.dtype = DataType::FP16,
@@ -280,7 +276,7 @@ static constexpr FmhaBwdDQDKDVVariant ALL_DQDKDV_VARIANTS[] = {
          .signature = {.dtype = DataType::FP16,
                        .hdim_q = 128, .hdim_v = 128,
                        .mode = FmhaMode::BATCH},
-         .algorithm = {.has_mask = true,
+         .algorithm = {.mask_type = FmhaMaskType::TOP_LEFT_CAUSAL,
                        .is_deterministic = true,
                        .pad_hdim_q = 8, .pad_hdim_v = 8}})},
 };
@@ -303,7 +299,7 @@ constexpr const FmhaBwdDQDKDVVariant* findVariant(FmhaBwdDQDKDVConfig cfg)
             continue;
 
         // Feature flags must match exactly
-        if(v.spec.has_mask != algo.has_mask || v.spec.has_dropout != algo.has_dropout ||
+        if(v.spec.mask_type != algo.mask_type || v.spec.has_dropout != algo.has_dropout ||
            v.spec.is_deterministic != algo.is_deterministic || v.spec.bias_type != algo.bias_type ||
            v.spec.has_bias_grad != algo.has_bias_grad)
             continue;

@@ -39,4 +39,27 @@ enum class FmhaBiasType
     ALIBI
 };
 
+/// Attention mask family.
+///
+/// Integer values must match ck_tile::GenericAttentionMaskEnum so the device
+/// bridge can forward the spec-time enum to the kernel via a static_cast
+/// without a translation table (see dqdkdv_dev.hpp's MASK_TYPE scalar wiring).
+///
+/// Both causal variants describe a sliding window with left=-1 (unbounded
+/// past) and right=0 (no lookahead). The two flavours differ in where the
+/// causal diagonal is anchored when seqlen_q != seqlen_k:
+///   * TOP_LEFT     -- diagonal at (0, 0); standard "predict next token".
+///   * BOTTOM_RIGHT -- diagonal at (seqlen_q-1, seqlen_k-1); used when the
+///                     query is the *tail* of a longer cached K/V (decode).
+/// GENERIC selects ck_tile's runtime (left, right, top-left/bottom-right)
+/// window descriptor and is intended for sliding-window / xformer-style
+/// attention. NO_MASK disables masking at compile time.
+enum class FmhaMaskType
+{
+    NO_MASK             = 0,
+    TOP_LEFT_CAUSAL     = 1,
+    BOTTOM_RIGHT_CAUSAL = 2,
+    GENERIC             = 3,
+};
+
 } // namespace rocm_ck
@@ -155,9 +155,10 @@ struct FmhaBwdDQDKDVTypes
                                                 kTile.max_seq_q>;
 
     // --- Mask type ---
-    // has_mask=true  -> GenericAttentionMask<true, true> (full masking)
-    // has_mask=false -> GenericAttentionMask<false>      (no masking)
-    using Mask = std::conditional_t<K.has_mask,
+    // mask_type==NO_MASK -> GenericAttentionMask<false>      (no masking)
+    // mask_type!=NO_MASK -> GenericAttentionMask<true, true> (full masking;
+    //                       runtime window/anchor selected via MASK_TYPE slot)
+    using Mask = std::conditional_t<hasMask(K),
                                     ck_tile::GenericAttentionMask<true, true>,
                                     ck_tile::GenericAttentionMask<false>>;
 
@@ -481,7 +482,7 @@ __device__ void runFmhaBwdDQDKDV(Args args)
         }
     }
 
-    if constexpr(K.has_mask)
+    if constexpr(hasMask(K))
     {
         kargs.window_size_left  = args.scalars[S::WINDOW_SIZE_LEFT].i32;
         kargs.window_size_right = args.scalars[S::WINDOW_SIZE_RIGHT].i32;

@@ -322,7 +322,7 @@ struct FmhaBwdDQDKDVAlgorithm
     // Feature flags -- which variation of the computation
     FmhaBiasType bias_type = FmhaBiasType::NONE;
     bool has_bias_grad     = false;
-    bool has_mask          = false;
+    FmhaMaskType mask_type = FmhaMaskType::NO_MASK;
     bool has_dropout       = false;
     bool is_deterministic  = false;
 
@@ -352,7 +352,7 @@ struct FmhaBwdDQDKDVSpec
     // From Algorithm -- feature flags
     FmhaBiasType bias_type;
     bool has_bias_grad;
-    bool has_mask;
+    FmhaMaskType mask_type;
     bool has_dropout;
     bool is_deterministic;
 
@@ -431,7 +431,7 @@ constexpr int P_UNDROP       = 4; // f32: 1/(1-dropout_rate), passed to CK Tile
 constexpr int RP_UNDROP      = 5; // f32: 1-dropout_rate (keep_prob), used for p_undrop_in_uint8_t
 constexpr int DROP_SEED      = 6; // u64: dropout RNG seed
 constexpr int DROP_OFFSET    = 7; // u64: dropout RNG offset
-// Mask scalar slots -- present only when has_mask=true.
+// Mask scalar slots -- present only when mask_type != NO_MASK.
 // Indices are fixed regardless of dropout; unused slots are not populated.
 constexpr int WINDOW_SIZE_LEFT  = 8;  // i32: left context window (-1 = unlimited)
 constexpr int WINDOW_SIZE_RIGHT = 9;  // i32: right context window (0 = causal)
@@ -445,6 +445,12 @@ constexpr int BATCH_SIZE = 11; // i32: batch count (deterministic batch mode onl
 
 } // namespace fmha_bwd_dqdkdv_slots
 
+/// True iff the spec enables any attention mask.
+/// Single source of truth for "is masking active" so the device bridge,
+/// host runner, registry matcher, and slot-count predicate cannot drift
+/// apart when a new mask family is added.
+constexpr bool hasMask(FmhaBwdDQDKDVSpec k) { return k.mask_type != FmhaMaskType::NO_MASK; }
+
 /// Single source of truth for "does this spec use the BATCH_SIZE scalar slot".
 /// Used by requiredScalars(), validateArgs(), and the device bridge so the
 /// predicate cannot drift between sites. CK Tile's persistent kernel
@@ -462,7 +468,7 @@ constexpr int requiredScalars(FmhaBwdDQDKDVSpec k)
 {
     if(usesBatchSizeSlot(k))
         return BATCH_SIZE + 1; // 12 (dominates mask/dropout slots)
-    if(k.has_mask)
+    if(hasMask(k))
         return MASK_TYPE + 1; // 11 (covers dropout slots [4..7] since 11 > 8)
     if(k.has_dropout)
         return DROP_OFFSET + 1; // 8
@@ -517,6 +523,17 @@ consteval FmhaBwdDQDKDVSpec makeSpec(FmhaBwdDQDKDVConfig cfg)
     if(algo.has_bias_grad && algo.bias_type == FmhaBiasType::NONE)
         throw "has_bias_grad requires bias_type != NONE";
 
+    // mask_type must be one of the four declared enum values. The cast guards
+    // against callers passing an integer-via-enum out of range -- the device
+    // bridge static_casts straight to ck_tile::GenericAttentionMaskEnum and an
+    // unknown value would silently land on undefined kernel behaviour.
+    {
+        const auto m = static_cast<int>(algo.mask_type);
+        if(m < static_cast<int>(FmhaMaskType::NO_MASK) ||
+           m > static_cast<int>(FmhaMaskType::GENERIC))
+            throw "mask_type must be NO_MASK, TOP_LEFT_CAUSAL, BOTTOM_RIGHT_CAUSAL, or GENERIC";
+    }
+
     // --- padding validation ---
     if(algo.pad_hdim_q != 0 && algo.pad_hdim_q != 1 && algo.pad_hdim_q != 8)
         throw "pad_hdim_q must be 0, 1, or 8";
@@ -547,7 +564,7 @@ consteval FmhaBwdDQDKDVSpec makeSpec(FmhaBwdDQDKDVConfig cfg)
         .mode             = sig.mode,
         .bias_type        = algo.bias_type,
         .has_bias_grad    = algo.has_bias_grad,
-        .has_mask         = algo.has_mask,
+        .mask_type        = algo.mask_type,
         .has_dropout      = algo.has_dropout,
         .is_deterministic = algo.is_deterministic,
         .pad_hdim_q       = algo.pad_hdim_q,
@@ -603,7 +620,7 @@ static_assert(fmha_bwd_dqdkdv_slots::requiredScalars(makeSpec(
         .signature = {.dtype = DataType::FP16,
                       .hdim_q = 128, .hdim_v = 128,
                       .mode = FmhaMode::BATCH},
-        .algorithm = {.has_mask = true,
+        .algorithm = {.mask_type = FmhaMaskType::TOP_LEFT_CAUSAL,
                       .is_deterministic = true,
                       .pad_hdim_q = 8, .pad_hdim_v = 8}})) == 12);
 

@@ -0,0 +1,19 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+//
+// Must fail: mask_type must be one of the four declared enum values.
+// A value cast in from out-of-range integer space (e.g. 99) is forwarded
+// verbatim to the device-side static_cast<ck_tile::GenericAttentionMaskEnum>,
+// where it would silently land on undefined kernel behaviour. The consteval
+// validator in makeSpec rejects it at compile time instead.
+//
+// Expected error: "mask_type must be NO_MASK, TOP_LEFT_CAUSAL,
+//                  BOTTOM_RIGHT_CAUSAL, or GENERIC"
+
+#include <rocm_ck/ops/fmha_bwd/dqdkdv_spec.hpp>
+
+using namespace rocm_ck;
+
+constexpr auto bad = makeSpec(FmhaBwdDQDKDVConfig{
+    .signature = {.dtype = DataType::FP16, .hdim_q = 128, .hdim_v = 128, .mode = FmhaMode::BATCH},
+    .algorithm = {.mask_type = static_cast<FmhaMaskType>(99), .pad_hdim_q = 8, .pad_hdim_v = 8}});