fix(hip-kernel-provider): ASM SDPA backward rejects the a16 accumulator

dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/asm/SdpaBwdKernelArgs.hpp

@@ -248,8 +248,8 @@ struct __attribute__((packed)) fmha_bwd_dqdkdv_args
 // Kernel 3: DQ_CONVERT Post-Processing Kernel
 // =============================================================================
 // Converts dQ from FP32 accumulator (dq_acc) to BF16 output format.
-// Only used with a32 accumulator kernels.  The a16 accumulator kernels write
-// dQ directly in BF16 and skip this step.
+// Only used with a32 accumulator kernels.  The a16 accumulator kernels add dQ
+// atomically into the output dQ, in its 16-bit type, and skip this step.
 //
 // Inputs:  dq_acc [B, H_q, S_q, D_qk] (FP32 from dqdkdv kernel)
 // Outputs: dQ [B, H_q, S_q, D_qk] (BF16)

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

@@ -14,12 +14,14 @@ namespace asm_sdpa_engine
 ///
 /// - A32: dQ is accumulated in FP32 into a workspace buffer (dq_acc), then a
 ///   separate DQ_CONVERT kernel converts FP32 → BF16. Requires 3 kernels.
-/// - A16: dQ is written directly in BF16 by the DQDKDV kernel. No dq_acc
-///   workspace buffer, no DQ_CONVERT kernel. Requires 2 kernels.
+/// - A16: dQ is accumulated directly in the output's 16-bit type: the DQDKDV
+///   kernel adds it atomically into the output dQ, which must be zeroed first.
+///   No dq_acc workspace buffer, no DQ_CONVERT kernel. Requires 2 kernels. Not
+///   supported: SdpaBwdPlanBuilder rejects it (see initializeExecutionSettings).
 enum class AccumulatorType : uint8_t
 {
     A32, // FP32 accumulator — 3-kernel path (ODO → DQDKDV → DQ_CONVERT)
-    A16 // BF16 accumulator — 2-kernel path (ODO → DQDKDV)
+    A16 // 16-bit accumulator — 2-kernel path (ODO → DQDKDV)
 };
 
 /**

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

@@ -173,7 +173,8 @@ asm_sdpa_engine::fmha_bwd_dqdkdv_args
     asm_sdpa_engine::fmha_bwd_dqdkdv_args dqdkdv{};
 
     // A32: write dQ to FP32 dq_acc workspace (DQ_CONVERT casts it to BF16 afterward).
-    // A16: write dQ directly to the output BF16 buffer.
+    // A16: add dQ atomically into the output dQ, which must be zeroed first (not
+    // supported; see SdpaBwdPlanBuilder::initializeExecutionSettings).
     dqdkdv.ptr_dq = (a.accType == asm_sdpa_engine::AccumulatorType::A32) ? a.dq_acc_ptr : a.dq_ptr;
     dqdkdv.ptr_dk = a.dk_ptr;
     dqdkdv.ptr_dv = a.dv_ptr;
@@ -488,13 +489,13 @@ void SdpaBwdPlan::execute(const Handle& handle,
 
     // 4. Carve workspace into sub-buffers.
     // A32: dq_acc follows D buffer (DQDKDV accumulates FP32 dQ there, then DQ_CONVERT casts).
-    // A16: DQDKDV writes dQ directly to the output buffer; dq_acc is not allocated.
+    // A16: DQDKDV adds dQ atomically into the output dQ; dq_acc is not allocated.
     // Layout: D buffer | dq_acc (A32 only)
     auto* dBufPtr = workspace;
     size_t wsOffset = sdpaBwdDBufferSize(_params.batchSize, _params.numHeadsQ, _params.seqLenQ);
 
     // A32: dq_acc buffer follows D buffer in workspace.
-    // A16: no dq_acc buffer (nullptr) — DQDKDV writes dQ directly to user output.
+    // A16: no dq_acc buffer (nullptr) — DQDKDV adds dQ atomically into the user's dQ.
     void* dqAccPtr = nullptr;
     if(_params.accumulatorType == AccumulatorType::A32)
     {
@@ -526,7 +527,7 @@ void SdpaBwdPlan::execute(const Handle& handle,
 
     // 7. Launch kernels on the same stream.
     // a32: 3 kernels — ODO → DQDKDV → DQ_CONVERT (sequential dependencies)
-    // a16: 2 kernels — ODO → DQDKDV (dQ written directly in BF16)
+    // a16: 2 kernels — ODO → DQDKDV (dQ added atomically into the output dQ)
     // Launching on the same stream guarantees ordering without explicit barriers.
     auto stream = handle.getStream();
 
@@ -560,7 +561,9 @@ void SdpaBwdPlan::execute(const Handle& handle,
     // dQ contributions atomically and does not pre-zero; stale residue from a
     // prior workspace lease would silently corrupt dQ. AITER allocates dq_accum
     // via torch::zeros (aiter/csrc/py_itfs_cu/asm_mha_bwd.cu:137 at commit 17d4a33).
-    // A16 writes dQ directly — no accumulator buffer needed, skip the memset.
+    // A16 never gets here (SdpaBwdPlanBuilder rejects it). Its kernels also add
+    // dQ atomically, into the output dQ, so enabling A16 means zeroing dQ here
+    // (AITER calls dq.zero_() first).
     if(_params.accumulatorType == AccumulatorType::A32)
     {
         const size_t dqAccBytes = sdpaBwdDqAccBufferSize(
@@ -592,7 +595,8 @@ void SdpaBwdPlan::execute(const Handle& handle,
     }
 
     // 7c. DQ_CONVERT (FP32 → BF16) — A32 path only.
-    // A16 wrote dQ directly to the output BF16 buffer in step 7b; no cast needed.
+    // A16 adds dQ atomically into the output dQ in step 7b, so it needs no cast,
+    // only the zeroing described there.
     if(_params.accumulatorType == AccumulatorType::A32)
     {
         auto postArgs = buildPostArgs(mhaArgs);

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

@@ -21,12 +21,15 @@ namespace asm_sdpa_engine
  *
  * Orchestrates 2 or 3 ASM kernels for the backward pass:
  *   1. ODO        — D reduction: D[b,h,i] = sum_j(O * dO)
- *   2. DQDKDV     — Main gradients: dQ (FP32 via A32, or direct BF16 via A16), dK, dV
+ *   2. DQDKDV     — Main gradients: dQ (FP32 into dq_acc via A32, or added
+ *                   atomically into the output dQ via A16), dK, dV
  *   3. DQ_CONVERT — Post-processing: FP32 dQ → BF16 (A32 path only)
  *
  * Accumulator type selects the path:
  *   - A32 (3-kernel path): ODO → DQDKDV → DQ_CONVERT
- *   - A16 (2-kernel path): ODO → DQDKDV (dQ written directly in BF16)
+ *   - A16 (2-kernel path): ODO → DQDKDV (dQ added atomically into the output
+ *     dQ, which must be zeroed first). Not supported: SdpaBwdPlanBuilder never
+ *     builds an A16 plan.
  *
  * When A16, `_postKernel` is std::nullopt and the DQ_CONVERT launch + dq_acc
  * workspace are both skipped.

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

@@ -306,7 +306,7 @@ bool wouldBwdByteStridesFitUint32(
 // and odo additionally always uses the bf16_cvt=3 sentinel. The dqdkdv
 // (main backward) kernel carries the full dispatch axes; its atomic32 column
 // follows the resolved accumulator type (A32 → 3-kernel FP32-accumulator path,
-// A16 → 2-kernel BF16 path), while pssk=1, pddv=1 are pinned because the
+// A16 → 2-kernel 16-bit path), while pssk=1, pddv=1 are pinned because the
 // registry rows for pssk=1, pddv=0 do not exist in batch mode and the unpadded
 // (pssk=0, pddv=0) row uses a different kernarg layout than the engine builds
 // today. TODO: lift the padding pin once the engine emits the unpadded kernarg
@@ -455,7 +455,13 @@ std::string lookupKernelNameKey(PipelineStage stage,
 
 } // namespace bwd_dispatch
 
-// Engine knob identifier for accumulator precision (a32 vs a16)
+// Engine knob identifier for accumulator precision. Only "a32" is supported;
+// initializeExecutionSettings explains why "a16" is rejected.
+//
+// AsmSdpaEngine publishes no builder knobs and never calls
+// initializeExecutionSettings, so plans built through hipDNN always use the A32
+// default. The knob is kept so the builder contract stays correct if the engine
+// ever forwards knobs the way HipMlopsEngine does.
 static constexpr const char* K_ACC_TYPE_KNOB_NAME = "sdpa.bwd.accumulator_type";
 
 bool SdpaBwdPlanBuilder::isApplicable(
@@ -682,9 +688,9 @@ bool SdpaBwdPlanBuilder::isApplicable(
 
     const int bf16CvtValue = (dataTypeId == "fp16") ? BF16_CVT_FP16_SENTINEL
                                                     : static_cast<int>(getRoundingMode(attrs));
-    // Applicability validates the default (A32) accumulator path: it is the
-    // default knob value and the most demanding (dq_acc + DQ_CONVERT). The user
-    // may still select A16 via the knob at buildPlan time.
+    // Applicability validates the A32 accumulator path, the only one this
+    // builder dispatches: initializeExecutionSettings rejects "a16" and
+    // buildPlan refuses an A16 setting.
     auto dispatchTuples = computeDispatchTuples(
         dataTypeId, headDimQk, maskType, bf16CvtValue, AccumulatorMode::A32);
 
@@ -783,6 +789,8 @@ size_t SdpaBwdPlanBuilder::getMaxWorkspaceSize(
     auto seqLenQ = static_cast<size_t>(qTensor->dims()->Get(2));
     auto headDimQk = static_cast<size_t>(qTensor->dims()->Get(3));
 
+    // An A16 setting is sized without dq_acc and not refused here. buildPlan is
+    // the gate: it refuses A16, so no plan ever runs with this size.
     const AccumulatorType accType
         = executionSettings.accumulatorType.value_or(AccumulatorType::A32);
 
@@ -815,17 +823,41 @@ void SdpaBwdPlanBuilder::initializeExecutionSettings(
     if(value == "a32")
     {
         executionSettings.accumulatorType = AccumulatorType::A32;
+        return;
     }
-    else if(value == "a16")
-    {
-        executionSettings.accumulatorType = AccumulatorType::A16;
-    }
-    else
+
+    // "a16" is rejected. Every DQDKDV dispatch tuple requests pssk=1, pddv=1,
+    // and gfx942 ships no a16 kernel for that layout, so buildPlan would fail
+    // for every graph isApplicable accepts. (gfx950, which isApplicable does
+    // not accept yet, has such rows, but none has been validated here.) The
+    // gfx942 a16 kernels that do ship (pssk=0) are no drop-in replacement
+    // either. They add dQ atomically, in the 16-bit output type (BF16 or FP16),
+    // straight into the output dQ, address dQ and dO with Q's strides, and
+    // never load seqlen_k, the window bounds, or any V, dO, dK or dV stride
+    // except dK's seq stride. isApplicable cannot do the gating: it sees no
+    // knob, and it accepts graphs for A32 that A16 cannot run. Enabling A16
+    // needs, besides a kernel row, getCustomKnobs to list "a16" and this
+    // function to accept it only for graphs where:
+    //   - Sq == Skv and Sq % 64 == 0;
+    //   - the mask is none or causal, not a sliding window;
+    //   - the head dim is 64 or 128;
+    //   - dQ and dO have Q's strides;
+    //   - V has K's strides, dK and dV have K's head and batch strides, and dV
+    //     has dK's seq stride;
+    //   - dQ can be zeroed in one memset (e.g. packed), which the plan must do
+    //     before DQDKDV.
+    // The a16 dQ also differs between runs on identical inputs.
+    if(value == "a16")
     {
-        throw hipdnn_plugin_sdk::HipdnnPluginException(HIPDNN_PLUGIN_STATUS_INVALID_VALUE,
-                                                       "Invalid accumulator_type: '" + value
-                                                           + "'. Must be 'a32' or 'a16'");
+        throw hipdnn_plugin_sdk::HipdnnPluginException(
+            HIPDNN_PLUGIN_STATUS_INVALID_VALUE,
+            "accumulator_type 'a16' is not supported by the ASM SDPA backward engine; "
+            "use 'a32'");
     }
+
+    throw hipdnn_plugin_sdk::HipdnnPluginException(HIPDNN_PLUGIN_STATUS_INVALID_VALUE,
+                                                   "Invalid accumulator_type: '" + value
+                                                       + "'. Must be 'a32'");
 }
 
 void SdpaBwdPlanBuilder::buildPlan(
@@ -835,13 +867,23 @@ void SdpaBwdPlanBuilder::buildPlan(
     Context& executionContext) const
 {
     // -------------------------------------------------------------------------
-    // 0. Resolve accumulator type from the user knob / execution settings.
-    // This is the single source of truth for a32 (3-kernel) vs a16 (2-kernel)
-    // selection; it drives CSV row selection, DQ_CONVERT loading, the dq_acc
-    // workspace, and params below. Defaults to A32 when the knob is unset.
+    // 0. Resolve accumulator type from the execution settings, defaulting to
+    // A32 when the knob is unset. It drives CSV row selection, DQ_CONVERT
+    // loading, the dq_acc workspace, and params below. Only A32 is supported:
+    // initializeExecutionSettings rejects "a16", and an A16 setting that gets
+    // here by any other route is refused before any kernel is resolved. This is
+    // the only gate: getMaxWorkspaceSize still sizes an A16 setting, and the
+    // A16 branches below are kept for a future enablement.
     // -------------------------------------------------------------------------
     const AccumulatorType accType
         = executionContext.executionSettings().accumulatorType.value_or(AccumulatorType::A32);
+    if(accType != AccumulatorType::A32)
+    {
+        throw hipdnn_plugin_sdk::HipdnnPluginException(
+            HIPDNN_PLUGIN_STATUS_INTERNAL_ERROR,
+            "SdpaBwdPlanBuilder::buildPlan: the A16 accumulator is not supported "
+            "(initializeExecutionSettings should have rejected it)");
+    }
 
     auto deviceStringOpt = plan_utils::tryGetDeviceString(
         handle.getStream(), "Failed to query device properties with error: ");
@@ -1003,9 +1045,9 @@ void SdpaBwdPlanBuilder::buildPlan(
     auto batchMode = getBatchMode(sdpaAttrs);
     const int bf16CvtValue = (dataTypeId == "fp16") ? BF16_CVT_FP16_SENTINEL
                                                     : static_cast<int>(getRoundingMode(sdpaAttrs));
-    // The resolved accumulator type (from the user knob, defaulting to A32) is
-    // authoritative: it selects the a32 vs a16 dqdkdv CSV row and gates whether
-    // the DQ_CONVERT stage is resolved and loaded below.
+    // accType is A32 here: step 0 refuses A16. It selects the dqdkdv CSV row's
+    // atomic32 value and gates whether the DQ_CONVERT stage is resolved and
+    // loaded below.
     auto dispatchTuples = computeDispatchTuples(dataTypeId,
                                                 static_cast<int>(headDimQk),
                                                 maskType,
@@ -1081,11 +1123,12 @@ void SdpaBwdPlanBuilder::buildPlan(
                                                   static_cast<int>(batchMode),
                                                   dqdtuple.bf16Cvt));
 
-    // The DQ_CONVERT stage (FP32 dq_acc → BF16 cast) is only part of the A32
-    // 3-kernel path. For A16 the DQDKDV kernel writes dQ directly in BF16, so
-    // dq_convert is not resolved, the dq_acc workspace is skipped, and the plan
-    // is built with a nullopt post-kernel. The resolved tuple's atomic32 column
-    // already mirrors accType (see computeDispatchTuples above).
+    // The DQ_CONVERT stage (FP32 dq_acc → output dtype cast) is only part of the
+    // A32 3-kernel path, the only path that gets here. For A16 (refused in step
+    // 0) the DQDKDV kernel adds dQ atomically into the output dQ, so dq_convert
+    // would not be resolved, the dq_acc workspace would be skipped, and the plan
+    // would be built with a nullopt post-kernel. The resolved tuple's atomic32
+    // column already mirrors accType (see computeDispatchTuples above).
     const bool useA32 = (accType == AccumulatorType::A32);
 
     std::optional<ResolvedKernel> dqConvertResolved;
@@ -1252,16 +1295,16 @@ std::vector<hipdnn_flatbuffers_sdk::data_objects::KnobT> SdpaBwdPlanBuilder::get
     KnobT accKnob;
     accKnob.knob_id = K_ACC_TYPE_KNOB_NAME;
     accKnob.description
-        = "Accumulator precision for backward dQ gradient: a32 (FP32, 3-kernel) or a16 (BF16, "
-          "2-kernel)";
+        = "Accumulator precision for backward dQ gradient. Only a32 (FP32, 3-kernel) is "
+          "supported";
 
     StringValueT defaultValue;
     defaultValue.value = "a32";
     accKnob.default_value.Set(defaultValue);
 
+    // "a16" is not listed: initializeExecutionSettings rejects it.
     StringConstraintT constraint;
     constraint.valid_values.emplace_back("a32");
-    constraint.valid_values.emplace_back("a16");
     accKnob.constraint.Set(constraint);
 
     knobs.push_back(std::move(accKnob));

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

@@ -57,7 +57,8 @@ constexpr size_t sdpaBwdDBufferSize(size_t batch, size_t headsQ, size_t seqLenQ)
 }
 
 /// dq_acc buffer: FP32 accumulator for dQ, shape [B, H_q, S_q, D_qk] in FP32.
-/// Only needed for a32-accumulator kernels; a16 kernels write dQ in BF16 directly.
+/// Only needed for a32-accumulator kernels; a16 kernels add dQ atomically into
+/// the output dQ.
 constexpr size_t sdpaBwdDqAccBufferSize(size_t batch, size_t headsQ, size_t seqLenQ, size_t headDim)
 {
     return alignUp(batch * headsQ * seqLenQ * headDim * sizeof(float), K_WORKSPACE_ALIGNMENT_BYTES);