fix(hipdnn): Set Sdpa diagonal alignment and left/right_bounds correctly

dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaBwdPlanBuilder.cpp

@@ -630,10 +630,14 @@ bool SdpaBwdPlanBuilder::isApplicable(
 
     // Classify the mask; contradictory mask attributes are an invalid-input
     // condition the engine declines rather than dispatches.
+    const int64_t seqLenQ = qTensor->dims()->Get(2);
+    const int64_t seqLenKv = kTensor->dims()->Get(2);
     MaskType maskType = MaskType::NO_MASK;
     try
     {
-        maskType = plan_utils::getMaskType(attrs);
+        const auto band = plan_utils::makeCoveringBoundsUnbound(
+            plan_utils::resolveDiagonalBand(attrs), seqLenQ, seqLenKv);
+        maskType = plan_utils::getMaskType(band.leftBound, band.rightBound, band.alignment);
     }
     catch(const hipdnn_plugin_sdk::HipdnnPluginException& e)
     {
@@ -712,7 +716,6 @@ bool SdpaBwdPlanBuilder::isApplicable(
     HIP_KERNEL_RETURN_FALSE_IF(!dqdkdvCfgOpt,
                                "Failed to resolve dqdkdv config for byte-stride validation");
     const bool useA32 = (dqdkdvTuple.atomic32 == static_cast<int>(AccumulatorMode::A32));
-    const int64_t seqLenQ = qTensor->dims()->Get(2);
     HIP_KERNEL_RETURN_FALSE_IF(!wouldBwdByteStridesFitUint32(*qTensor,
                                                              *kTensor,
                                                              *vTensor,
@@ -959,7 +962,10 @@ void SdpaBwdPlanBuilder::buildPlan(
             "(isApplicable should have rejected)");
     }
     const auto& dataTypeId = *dataTypeIdOpt;
-    auto maskType = plan_utils::getMaskType(sdpaAttrs);
+    const auto maskBand = plan_utils::makeCoveringBoundsUnbound(
+        plan_utils::resolveDiagonalBand(sdpaAttrs), seqLenQ, seqLenKv);
+    const auto maskType
+        = plan_utils::getMaskType(maskBand.leftBound, maskBand.rightBound, maskBand.alignment);
     auto batchMode = getBatchMode(sdpaAttrs);
     const int bf16CvtValue = (dataTypeId == "fp16") ? BF16_CVT_FP16_SENTINEL
                                                     : static_cast<int>(getRoundingMode(sdpaAttrs));
@@ -1174,14 +1180,10 @@ void SdpaBwdPlanBuilder::buildPlan(
     params.maskOrdinal = static_cast<int32_t>(maskType);
     if(maskType == MaskType::SLIDING_WINDOW)
     {
-        params.windowLeft = sdpaAttrs.left_bound().has_value()
-                                ? static_cast<int32_t>(sdpaAttrs.left_bound().value())
-                                : -1;
-        params.windowRight = sdpaAttrs.right_bound().has_value()
-                                 ? static_cast<int32_t>(sdpaAttrs.right_bound().value())
-                                 : -1;
+        params.windowLeft = static_cast<int32_t>(maskBand.leftBound);
+        params.windowRight = static_cast<int32_t>(maskBand.rightBound);
         params.topLeftAlignment
-            = sdpaAttrs.diagonal_alignment()
+            = maskBand.alignment
               != hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment::BOTTOM_RIGHT;
     }
 

dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdPlanBuilder.cpp

@@ -445,7 +445,11 @@ bool SdpaFwdPlanBuilder::isApplicable(
     plan_utils::MaskType maskType = plan_utils::MaskType::NO_MASK;
     try
     {
-        maskType = plan_utils::getMaskType(attrs);
+        const auto band
+            = plan_utils::makeCoveringBoundsUnbound(plan_utils::resolveDiagonalBand(attrs),
+                                                    qTensor->dims()->Get(2),
+                                                    kTensor->dims()->Get(2));
+        maskType = plan_utils::getMaskType(band.leftBound, band.rightBound, band.alignment);
     }
     catch(const hipdnn_plugin_sdk::HipdnnPluginException& e)
     {
@@ -645,7 +649,9 @@ void SdpaFwdPlanBuilder::buildPlan(
     params.lseStrideHead = lseStrideHead;
     params.attnScale = attnScale;
     params.archString = deviceString;
-    params.maskType = plan_utils::getMaskType(sdpaAttrs);
+    const auto band = plan_utils::makeCoveringBoundsUnbound(
+        plan_utils::resolveDiagonalBand(sdpaAttrs), seqLenQ, seqLenKv);
+    params.maskType = plan_utils::getMaskType(band.leftBound, band.rightBound, band.alignment);
 
     const auto dataTypeId = getDataTypeIdentifier(
         qTensor->data_type(), kTensor->data_type(), vTensor->data_type(), oTensor->data_type());

dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaPlanUtils.hpp

@@ -55,35 +55,61 @@ enum class MaskType : int
     SLIDING_WINDOW = 3
 };
 
-// Classify the mask requested by an SDPA (forward or backward) attribute set.
+// The diagonal band an SDPA (forward or backward) attribute set requests, with
+// the deprecated causal_mask / causal_mask_bottom_right booleans merged in.
+struct DiagonalBand
+{
+    int64_t leftBound;
+    int64_t rightBound;
+    hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment alignment;
+};
+
+// Replace each bound that already covers the whole sequence span with -1
+// (unbounded).
 //
-// Two sources can describe the mask: the modern left_bound / right_bound /
-// diagonal_alignment trio, and the deprecated causal_mask /
-// causal_mask_bottom_right booleans. When a deprecated boolean is set it takes
-// precedence and the modern trio is ignored; otherwise the trio is
-// authoritative. The two deprecated booleans are mutually exclusive, so setting
-// both throws HipdnnPluginException(INVALID_VALUE).
+// A bound that reaches the widest offset the band can span for this alignment
+// and these sequence lengths masks nothing on its side, so it becomes -1. That
+// is the same mask: computeMaskCoordinates() replaces -1 with exactly that
+// widest offset (seqLen - 1 on the matching axis). A covering window therefore
+// classifies as plain causal or no-mask. Bounds that survive are returned
+// unchanged as int64, so a kernel field filled from one must first confirm it
+// fits int32 (every survivor is below a sequence length, which does).
 //
-// Guaranteeing the two parameter sets agree belongs in the hipDNN frontend; this
-// helper only resolves which source wins for dispatch.
+// A sequence length of 0 or 1 gives a span <= 0, so a bound of 0 also becomes
+// -1 (e.g. top-left causal with seqLenKv == 1 is unmasked).
+//
+// Expects bounds >= -1 (resolveDiagonalBand() rejects the rest).
+inline DiagonalBand makeCoveringBoundsUnbound(DiagonalBand band, int64_t seqLenQ, int64_t seqLenKv)
+{
+    using hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment;
+    const bool bottomRight = band.alignment == DiagonalAlignment::BOTTOM_RIGHT;
+    const int64_t leftSpan = bottomRight ? seqLenKv - 1 : seqLenQ - 1;
+    const int64_t rightSpan = bottomRight ? seqLenQ - 1 : seqLenKv - 1;
+    const auto unboundIfCovering
+        = [](int64_t bound, int64_t span) { return (bound < 0 || bound >= span) ? -1 : bound; };
+    band.leftBound = unboundIfCovering(band.leftBound, leftSpan);
+    band.rightBound = unboundIfCovering(band.rightBound, rightSpan);
+    return band;
+}
+
+// Extract the requested diagonal band, merging the deprecated booleans the way
+// cuDNN's setters compose: a boolean acts as if its setter ran first (right
+// bound 0, plus BOTTOM_RIGHT alignment for causal_mask_bottom_right) and the
+// stored bounds and alignment are applied on top. An unset bound, or one equal
+// to -1, does not override the causal right bound. causal_mask keeps the stored
+// alignment; causal_mask_bottom_right always resolves to BOTTOM_RIGHT, because
+// TOP_LEFT is the schema default and indistinguishable from an unset alignment.
 //
-// Absence-awareness: the generated flatbuffer accessors expose the causal_mask*
-// fields as plain bool defaulting to false, with no has_*() accessor.
-// "Explicitly false" and "unset" are therefore indistinguishable; a false bool
-// is treated as "not requested". left_bound / right_bound are
-// flatbuffers::Optional, but an unset bound is treated as unbounded (-1) to
-// match the canonical convention used across the SDPA path, so a partially
-// specified trio (e.g. only right_bound = 0) still derives a mask rather than
-// silently falling back to NO_MASK.
+// Setting both booleans, or a bound below -1, throws
+// HipdnnPluginException(INVALID_VALUE).
 template <typename SdpaAttrsT>
-MaskType getMaskType(const SdpaAttrsT& attrs)
+DiagonalBand resolveDiagonalBand(const SdpaAttrsT& attrs)
 {
     using namespace hipdnn_flatbuffers_sdk::data_objects;
 
     const bool causalDeprecated = attrs.causal_mask();
     const bool bottomRightDeprecated = attrs.causal_mask_bottom_right();
 
-    // The two deprecated booleans are mutually exclusive.
     if(causalDeprecated && bottomRightDeprecated)
     {
         throw hipdnn_plugin_sdk::HipdnnPluginException(
@@ -92,31 +118,43 @@ MaskType getMaskType(const SdpaAttrsT& attrs)
             "but both are set");
     }
 
-    // Deprecated booleans take precedence: when either is set, defer to it and
-    // ignore the modern bounds trio.
-    if(causalDeprecated)
+    const int64_t left = attrs.left_bound().has_value() ? attrs.left_bound().value() : -1;
+    const int64_t right = attrs.right_bound().has_value() ? attrs.right_bound().value() : -1;
+
+    if(left < -1 || right < -1)
     {
-        return MaskType::TOP_LEFT_CAUSAL;
+        throw hipdnn_plugin_sdk::HipdnnPluginException(
+            HIPDNN_PLUGIN_STATUS_INVALID_VALUE,
+            "SDPA: left_bound and right_bound must be >= -1 (got left_bound=" + std::to_string(left)
+                + ", right_bound=" + std::to_string(right) + ")");
     }
-    if(bottomRightDeprecated)
+
+    if(!causalDeprecated && !bottomRightDeprecated)
     {
-        return MaskType::BOTTOM_RIGHT_CAUSAL;
+        return {left, right, attrs.diagonal_alignment()};
     }
 
-    // No deprecated boolean set: the modern bounds trio is authoritative. An
-    // unset bound means unbounded, represented here as -1, so a partially
-    // specified trio still resolves to the mask it describes.
-    const int64_t left = attrs.left_bound().has_value() ? attrs.left_bound().value() : -1;
-    const int64_t right = attrs.right_bound().has_value() ? attrs.right_bound().value() : -1;
-    if(left == -1 && right == -1) // both unbounded
+    return {left,
+            right >= 0 ? right : 0,
+            bottomRightDeprecated ? DiagonalAlignment::BOTTOM_RIGHT : attrs.diagonal_alignment()};
+}
+
+// Classify a diagonal band (from resolveDiagonalBand(), typically passed through
+// makeCoveringBoundsUnbound() first).
+inline MaskType getMaskType(int64_t leftBound,
+                            int64_t rightBound,
+                            hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment alignment)
+{
+    using hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment;
+
+    if(leftBound == -1 && rightBound == -1) // both unbounded
     {
         return MaskType::NO_MASK;
     }
-    if(left == -1 && right == 0) // causal: attend up to the diagonal
+    if(leftBound == -1 && rightBound == 0) // causal: attend up to the diagonal
     {
-        return attrs.diagonal_alignment() == DiagonalAlignment::BOTTOM_RIGHT
-                   ? MaskType::BOTTOM_RIGHT_CAUSAL
-                   : MaskType::TOP_LEFT_CAUSAL;
+        return alignment == DiagonalAlignment::BOTTOM_RIGHT ? MaskType::BOTTOM_RIGHT_CAUSAL
+                                                            : MaskType::TOP_LEFT_CAUSAL;
     }
     return MaskType::SLIDING_WINDOW; // anything else is a sliding window
 }

dnn-providers/hip-kernel-provider/src/engines/hip_flash2_engine/HipFlash2FwdPlanBuilder.cpp

@@ -71,14 +71,6 @@ bool HipFlash2FwdPlanBuilder::isApplicable(const Handle& handle,
     HIP_KERNEL_RETURN_FALSE_IF(attrs.page_table_v_tensor_uid(), "page_table_v not supported");
     HIP_KERNEL_RETURN_FALSE_IF(attrs.generate_stats().value_or(false),
                                "LSE stats output not supported");
-    // K2: reject mask types the kernel does not implement
-    {
-        const auto maskType = asm_sdpa_engine::plan_utils::getMaskType(attrs);
-        HIP_KERNEL_RETURN_FALSE_IF(
-            maskType != asm_sdpa_engine::plan_utils::MaskType::NO_MASK
-                && maskType != asm_sdpa_engine::plan_utils::MaskType::TOP_LEFT_CAUSAL,
-            "Only NO_MASK and TOP_LEFT_CAUSAL are supported");
-    }
     HIP_KERNEL_RETURN_FALSE_IF(attrs.seq_len_q_tensor_uid().has_value()
                                    || attrs.seq_len_kv_tensor_uid().has_value(),
                                "variable-length (group) batch mode not supported");
@@ -95,6 +87,29 @@ bool HipFlash2FwdPlanBuilder::isApplicable(const Handle& handle,
     HIP_KERNEL_RETURN_FALSE_IF(vTensor->dims()->size() != 4, "V must be rank-4");
     HIP_KERNEL_RETURN_FALSE_IF(oTensor->dims()->size() != 4, "O must be rank-4");
 
+    // K2: reject mask types the kernel does not implement
+    {
+        auto maskType = asm_sdpa_engine::plan_utils::MaskType::NO_MASK;
+        try
+        {
+            const auto band = asm_sdpa_engine::plan_utils::makeCoveringBoundsUnbound(
+                asm_sdpa_engine::plan_utils::resolveDiagonalBand(attrs),
+                qTensor->dims()->Get(2),
+                kTensor->dims()->Get(2));
+            maskType = asm_sdpa_engine::plan_utils::getMaskType(
+                band.leftBound, band.rightBound, band.alignment);
+        }
+        catch(const hipdnn_plugin_sdk::HipdnnPluginException& e)
+        {
+            HIPDNN_PLUGIN_LOG_INFO(std::string{HIP_KERNEL_LOG_PREFIX} + e.what());
+            return false;
+        }
+        HIP_KERNEL_RETURN_FALSE_IF(
+            maskType != asm_sdpa_engine::plan_utils::MaskType::NO_MASK
+                && maskType != asm_sdpa_engine::plan_utils::MaskType::TOP_LEFT_CAUSAL,
+            "Only NO_MASK and TOP_LEFT_CAUSAL are supported");
+    }
+
     // Data type: FP16 only
     const auto qType = qTensor->data_type();
     const auto kType = kTensor->data_type();
@@ -320,7 +335,7 @@ std::vector<data_objects::KnobT>
 }
 
 // ---------------------------------------------------------------------------
-// extractParams (private helper) -- Finding 1 fix: restored missing body
+// extractParams -- Finding 1 fix: restored missing body
 // ---------------------------------------------------------------------------
 Flash2FwdParams HipFlash2FwdPlanBuilder::extractParams(const Handle& /*handle*/,
                                                        const flatbuffer_utilities::IGraph& opGraph)
@@ -351,7 +366,13 @@ Flash2FwdParams HipFlash2FwdPlanBuilder::extractParams(const Handle& /*handle*/,
 
     p.attnScale = attnScaleFor(attrs);
 
-    p.causal = attrs.causal_mask();
+    const auto band = asm_sdpa_engine::plan_utils::makeCoveringBoundsUnbound(
+        asm_sdpa_engine::plan_utils::resolveDiagonalBand(attrs),
+        q->dims()->Get(2),
+        k->dims()->Get(2));
+    p.causal
+        = asm_sdpa_engine::plan_utils::getMaskType(band.leftBound, band.rightBound, band.alignment)
+          == asm_sdpa_engine::plan_utils::MaskType::TOP_LEFT_CAUSAL;
 
     p.qStrideBatch = q->strides()->Get(0);
     p.qStrideHead = q->strides()->Get(1);

dnn-providers/hip-kernel-provider/src/engines/hip_flash2_engine/HipFlash2FwdPlanBuilder_v2.hpp

@@ -77,7 +77,6 @@ class HipFlash2FwdPlanBuilder : public hipdnn_plugin_sdk::IPlanBuilder<Handle, S
         const Handle& handle,
         const hipdnn_flatbuffers_sdk::flatbuffer_utilities::IGraph& opGraph) const override;
 
-private:
     /// Extract Flash2FwdParams from a validated SDPA graph.
     static Flash2FwdParams
         extractParams(const Handle& handle,