diff --git a/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdParams.hpp b/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdParams.hpp index 060302cce009..2ec1f239a19a 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdParams.hpp +++ b/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdParams.hpp @@ -84,7 +84,9 @@ struct SdpaFwdParams unsigned int oStrideBatch; // LSE tensor stride (in elements). The forward kernel args struct - // (fmha_fwd_v3_args) only exposes s_lse_Hs — no batch stride field. + // (fmha_fwd_v3_args) only exposes s_lse_Hs — no batch stride field. The + // kernel steps between batches by numHeadsQ * s_lse_Hs, so isApplicable + // declines a stats tensor whose batch stride differs from that. unsigned int lseStrideHead = 0; // Tile size diff --git a/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdPlanBuilder.cpp b/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdPlanBuilder.cpp index 7f5d72aa700e..8eecd158a661 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdPlanBuilder.cpp +++ b/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaFwdPlanBuilder.cpp @@ -8,7 +8,9 @@ #include "plans/SdpaFwdPlan.hpp" #include "plans/SdpaPlanUtils.hpp" +#include #include +#include #include #include @@ -327,6 +329,152 @@ static bool wouldFwdByteStridesFitUint32( return ok; } +// Formats dims or strides for a decline reason. +static std::string vectorToString(const flatbuffers::Vector* values) +{ + return hipdnn_data_sdk::utilities::vecToString( + hipdnn_flatbuffers_sdk::utilities::convertFlatBufferVectorToStdVector(values)); +} + +// Q/K/V/O strides the kernels cannot address. Logs the first one found and +// returns false; returns true when every layout is addressable. Called by +// isApplicable after it has checked that Q, K, V and O have rank-4 dims. +// +// The forward argument block carries only batch, head and sequence byte strides +// for Q, K, V and O, and each row is read or written as D contiguous elements. On +// gfx942 each kind of layout declined here gave a wrong O in the cases run, +// except where noted: +// - Head-dim (innermost) stride other than 1: it is treated as 1 (a stride-2 Q, K +// or V read the gaps between elements, so O changed between runs; a stride-2 O +// left half of O unwritten). +// - Sequence stride below the tensor's own head dim, also 0 and also on a +// length-1 sequence axis: each buffer ends S * seq stride bytes after the start +// of its (batch, head) slice, so the end of the last row reads as 0 and writes +// there are dropped. (With S_kv = 1 a short K row still gave the right O, since +// the one key takes all the softmax weight; by the address formula its score, +// and so the stats, are wrong.) This relies on K having Q's head dim and O +// being [B, H_q, S_q, D_v], which isApplicable does not check yet: the kernels +// size K rows by Q's head dim and O's buffer by S_q. +// - O batch or head stride 0 on an axis longer than 1. Each (batch, head) slice +// of O is written by its own workgroups (grid z is the batch, and x or y the +// head), so such slices share memory and overwrite each other: an O batch or +// head stride of 0 left half or two thirds of O wrong, and O changed between +// some runs with identical inputs. +// Only a zero stride is checked, as in the backward: it is how a broadcast is +// written, so O gets one when it is given a broadcast input's strides. Nonzero +// strides that make O overlap itself, such as equal batch and head strides, +// have to be built by hand; they are not checked, although on gfx942 they also +// gave a wrong O. Stats do get an overlap check (below), because their address +// formula leaves a single comparison to make; for O it would take a footprint +// check over three axes in any order, which would fit better in hipDNN's graph +// validation, where every engine gets it. The batch-axis check applies only +// in batch mode: group-mode kernels locate each batch by its sequence offsets +// and ignore O's batch stride. +// A zero batch or head stride on Q, K or V (a broadcast) just re-reads one slice +// and stays accepted. So do odd strides, which the backward declines on some +// axes: with bf16 (hd128 and hd192) they were addressed correctly on every axis +// of Q, K, V and O, on gfx942 and on gfx950, whose kernels load rows as 16-byte +// vectors. The fp8 kernels, with their 1-byte elements, were not run with them. +static bool + areFwdIoLayoutsAddressable(const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes& q, + const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes& k, + const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes& v, + const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes& o, + BatchMode batchMode) +{ + // NOLINTNEXTLINE(readability-identifier-naming) + static const char* HIP_KERNEL_LOG_PREFIX = "[SdpaFwdPlanBuilder::isApplicable] "; + + struct IoTensor + { + const char* name; + const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes* tensor; + // Written by the kernels (O). + bool isOutput; + }; + const std::array ioTensors{ + {{"q", &q, false}, {"k", &k, false}, {"v", &v, false}, {"o", &o, true}}}; + for(const auto& io : ioTensors) + { + const auto* dims = io.tensor->dims(); + const auto* strides = io.tensor->strides(); + const auto describe = [&]() { + return std::string(io.name) + " tensor (Actual dims: " + vectorToString(dims) + + ", strides: " + vectorToString(strides) + ")"; + }; + HIP_KERNEL_RETURN_FALSE_IF(strides->size() != 4, + "rank-4 strides required for " + describe()); + HIP_KERNEL_RETURN_FALSE_IF(strides->Get(3) != 1, + "head-dim (innermost) stride must be 1 for " + describe()); + HIP_KERNEL_RETURN_FALSE_IF(strides->Get(2) < dims->Get(3), + "sequence stride must be at least the head dim for " + + describe()); + const bool zeroOutputStride + = io.isOutput + && ((batchMode == BatchMode::BATCH && dims->Get(0) > 1 && strides->Get(0) == 0) + || (dims->Get(1) > 1 && strides->Get(1) == 0)); + HIP_KERNEL_RETURN_FALSE_IF(zeroOutputStride, + "zero batch or head stride aliases output " + describe()); + } + return true; +} + +// Stats (LSE) layouts the kernels cannot write. Logs the first problem found and +// returns false; returns true when the layout is addressable. Called by +// isApplicable after it has checked that stats are [B, H_q, S_q] or +// [B, H_q, S_q, 1] matching Q. +// +// Every forward kernel writes the LSE of (b, h, s) at +// ((b * H_q + h) * lse_Hs + s) * 4 bytes, with lse_Hs the stats head stride +// (buildPlan); it takes no stats batch or sequence stride (from the .co +// disassembly of every forward kernel). On gfx942 a padded stats batch stride, a +// seq stride of 2, and [B, S_q, H] or [H, B, S_q] storage each gave wrong or NaN +// stats; by the address formula [H, B, S_q] also writes past the end of the stats +// buffer. So the seq stride must be 1 and the batch stride H_q times the head +// stride. A head stride below S_q makes (batch, head) slices overlap, so it +// declines too. Unlike in the backward, whose kernels ignore the stats strides, a +// padded head stride with a matching batch stride fits this formula, was correct +// on gfx942, and stays accepted. The stride of a size-1 batch or sequence axis is +// not checked. A size-1 head axis is, on purpose: buildPlan passes the head +// stride as lse_Hs, so with H_q = 1 and B > 1 the kernel steps between batches by +// the head stride, which must then equal the batch stride and be at least S_q. +// Passing the batch stride as lse_Hs when H_q = 1 would also accept packed +// [B, 1, S_q] stats with any head stride; such stats are rare and decline. +static bool + isFwdStatsLayoutAddressable(const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes& stats, + const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes& q) +{ + // NOLINTNEXTLINE(readability-identifier-naming) + static const char* HIP_KERNEL_LOG_PREFIX = "[SdpaFwdPlanBuilder::isApplicable] "; + + const int64_t batchSize = q.dims()->Get(0); + const int64_t numHeadsQ = q.dims()->Get(1); + const int64_t seqLenQ = q.dims()->Get(2); + const auto* statsStrides = stats.strides(); + const auto describe = [&]() { + return " (Actual stats dims: " + vectorToString(stats.dims()) + ", strides: " + + vectorToString(statsStrides) + ", q dims: " + vectorToString(q.dims()) + ")"; + }; + HIP_KERNEL_RETURN_FALSE_IF(statsStrides->size() != stats.dims()->size(), + "stats tensor must have one stride per dim" + describe()); + HIP_KERNEL_RETURN_FALSE_IF(seqLenQ > 1 && statsStrides->Get(2) != 1, + "stats tensor seq stride must be 1 when S_q > 1" + describe()); + // The kernel's step between batches; no match if it overflows int64. + int64_t kernelBatchStride = 0; + const bool batchStrideMatches + = !__builtin_mul_overflow(numHeadsQ, statsStrides->Get(1), &kernelBatchStride) + && statsStrides->Get(0) == kernelBatchStride; + HIP_KERNEL_RETURN_FALSE_IF(batchSize > 1 && !batchStrideMatches, + "stats tensor batch stride must be H_q times its head stride when " + "B > 1" + + describe()); + HIP_KERNEL_RETURN_FALSE_IF((batchSize > 1 || numHeadsQ > 1) && statsStrides->Get(1) < seqLenQ, + "stats tensor head stride must be at least S_q when B * H_q > 1 " + "(slices would overlap)" + + describe()); + return true; +} + static std::string getKernelTocKey(std::string coName, const std::string& archId, bool isMi308) { if(archId == "gfx942") @@ -480,6 +628,11 @@ bool SdpaFwdPlanBuilder::isApplicable( return false; } + if(!areFwdIoLayoutsAddressable(*qTensor, *kTensor, *vTensor, *oTensor, batchMode)) + { + return false; + } + // Validate optional stats (LSE) output tensor. // Gate purely on generate_stats; the UID is required when the flag is set. const bool hasStats = hasStatsOutput(attrs); @@ -525,6 +678,11 @@ bool SdpaFwdPlanBuilder::isApplicable( || statsTensor->dims()->Get(1) != qTensor->dims()->Get(1) || statsTensor->dims()->Get(2) != qTensor->dims()->Get(2), "stats tensor shape [B,H,Sq] must match Q tensor"); + + if(!isFwdStatsLayoutAddressable(*statsTensor, *qTensor)) + { + return false; + } } HIP_KERNEL_RETURN_FALSE_IF(qTensor->data_type() != kTensor->data_type() diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaFwdPlanBuilder.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaFwdPlanBuilder.cpp index f9063235089a..e25f994adfc0 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaFwdPlanBuilder.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaFwdPlanBuilder.cpp @@ -32,8 +32,15 @@ #include #include #include + +#include +#include +#include #include +#include +#include #include +#include namespace asm_sdpa_engine { @@ -375,6 +382,47 @@ auto createRaggedBshdSdpaFwdGraph(const std::vector& qDims = {4, 8, 256 return sdpa::withRaggedTensors(createBshdSdpaFwdGraph(qDims, vDims), {"q", "k", "v", "o"}); } +TEST_F(TestSdpaFwdPlanBuilder, IsApplicableIgnoresOnlyRaggedOutputBatchStride) +{ + using namespace hipdnn_flatbuffers_sdk::data_objects; + + SKIP_IF_NO_DEVICES(); + + const std::string deviceString + = hip_kernel_provider_common::getDeviceString(_handle.getStream()); + if(deviceString != "gfx942" && deviceString != "gfx950") + { + GTEST_SKIP(); + } + + const auto withOutputStrides = [](const std::vector& strides, bool ragged) { + auto builder + = mutateGraph(createBshdSdpaFwdGraph({2, 2, 256, 128}, {2, 2, 256, 128}), + [&](GraphT& graph) { findTensor(graph, "o").strides = strides; }); + if(ragged) + { + // Derive the offset multipliers from the overridden strides so each + // case satisfies the ragged-offset contract and isolates a layout rule. + return sdpa::withRaggedTensors(std::move(builder), {"q", "k", "v", "o"}); + } + return builder; + }; + + const std::vector> applicabilityTests = { + {GraphTest{withOutputStrides({0, 128, 256, 1}, false), "Dense O batch stride 0"}, false}, + {GraphTest{withOutputStrides({0, 128, 256, 1}, true), "Ragged O batch stride 0"}, true}, + {GraphTest{withOutputStrides({0, 0, 256, 1}, true), "Ragged O head stride 0"}, false}, + {GraphTest{withOutputStrides({0, 128, 256, 2}, true), "Ragged O head-dim stride 2"}, false}, + {GraphTest{withOutputStrides({0, 128, 64, 1}, true), "Ragged O sequence stride below D"}, + false}}; + + for(const auto& [test, applicable] : applicabilityTests) + { + EXPECT_EQ(_planBuilder.isApplicable(_handle, test.graphWrapper()), applicable) + << test.message; + } +} + TEST_F(TestSdpaFwdPlanBuilder, IsApplicableRaggedVariations) { using namespace hipdnn_flatbuffers_sdk::data_objects; @@ -2013,6 +2061,486 @@ TEST_F(TestSdpaFwdPlanBuilder, IsApplicableBoundsEachFp8VByteStrideField) } } +// ============================================================================= +// Tensor layouts the forward kernels can and cannot address +// ============================================================================= + +// Replaces the strides, and the dims when non-empty, of the tensor called name. +// createSdpaFwdGraph names its tensors q, k, v, o and stats. +struct LayoutOverride +{ + std::string name; + std::vector strides; + std::vector dims; +}; + +// Copy of a forward graph with the given tensor layouts applied. +flatbuffers::FlatBufferBuilder withLayouts(const flatbuffers::FlatBufferBuilder& base, + const std::vector& layouts) +{ + using namespace hipdnn_flatbuffers_sdk::data_objects; + + GraphT graph; + GetGraph(base.GetBufferPointer())->UnPackTo(&graph); + for(const auto& layout : layouts) + { + const auto it + = std::find_if(graph.tensors.begin(), graph.tensors.end(), [&](const auto& tensor) { + return tensor->name == layout.name; + }); + if(it == graph.tensors.end()) + { + throw std::logic_error("withLayouts: no tensor named " + layout.name); + } + (*it)->strides = layout.strides; + if(!layout.dims.empty()) + { + (*it)->dims = layout.dims; + } + } + + flatbuffers::FlatBufferBuilder builder; + builder.Finish(Graph::Pack(builder, &graph)); + return builder; +} + +constexpr std::array K_IO_TENSOR_NAMES = {"q", "k", "v", "o"}; + +// The same strides on the named tensors. +std::vector onTensors(const std::vector& names, + const std::vector& strides) +{ + std::vector layouts; + layouts.reserve(names.size()); + for(const char* name : names) + { + layouts.push_back({name, strides, {}}); + } + return layouts; +} + +// The same strides on all four Q/K/V/O tensors. +std::vector onAllIoTensors(const std::vector& strides) +{ + return onTensors({K_IO_TENSOR_NAMES.begin(), K_IO_TENSOR_NAMES.end()}, strides); +} + +// Packed bf16 forward graph with no mask. Q is qDims [B, H_q, S_q, D], V is vDims +// [B, H_kv, S_kv, D_v], K is [B, H_kv, S_kv, D] and O is [B, H_q, S_q, D_v]. With +// withStats, stats are packed [B, H_q, S_q, 1]. +flatbuffers::FlatBufferBuilder createLayoutGraph(const std::vector& qDims, + const std::vector& vDims, + bool withStats = false) +{ + using hipdnn_flatbuffers_sdk::data_objects::DataType; + return createSdpaFwdGraph(qDims, + vDims, + DataType::BFLOAT16, + DataType::FLOAT, + /*withAttnMask=*/false, + /*withScale=*/false, + withStats); +} + +// Layouts the kernels address correctly. These were GPU-verified with bf16 on +// gfx942: BSHD, padded and odd row strides, odd head and batch strides, broadcast +// (stride-0) inputs, and packed stats or stats with a padded head stride. The +// padded and odd strides were also verified on gfx950. The size-1-axis cases and +// rank-3 stats differ from them only in strides or dims no kernel reads. +TEST_F(TestSdpaFwdPlanBuilder, IsApplicableAcceptsAddressableLayouts) +{ + SKIP_IF_NO_DEVICES(); + + const std::string deviceString + = hip_kernel_provider_common::getDeviceString(_handle.getStream()); + if(deviceString != "gfx942" && deviceString != "gfx950") + { + GTEST_SKIP(); + } + + // B=2 and H=2 so that every batch and head stride is used. + const std::vector dims = {2, 2, 256, 128}; + // [B, H, S, D] dims stored as [B, S, H, D]. + const std::vector bshd = {int64_t{256} * 2 * 128, 128, int64_t{2} * 128, 1}; + // BHSD with each row padded from D=128 to D+8, D+2 or D+1 elements. + const std::vector rowsPlus8 = {int64_t{2} * 256 * 136, int64_t{256} * 136, 136, 1}; + const std::vector rowsPlus2 = {int64_t{2} * 256 * 130, int64_t{256} * 130, 130, 1}; + const std::vector rowsPlus1 = {int64_t{2} * 256 * 129, int64_t{256} * 129, 129, 1}; + // BHSD with an odd head or batch stride. + const std::vector oddHead + = {int64_t{2} * (int64_t{256} * 128 + 1), int64_t{256} * 128 + 1, 128, 1}; + const std::vector oddBatch = {int64_t{2} * 256 * 128 + 1, int64_t{256} * 128, 128, 1}; + // BHSD with the batch or head axis broadcast. + const std::vector zeroBatch = {0, int64_t{256} * 128, 128, 1}; + const std::vector zeroHead = {int64_t{256} * 128, 0, 128, 1}; + + std::vector> applicabilityTests; + applicabilityTests.emplace_back(GraphTest{createLayoutGraph(dims, dims), "packed BHSD"}, true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), onAllIoTensors(bshd)), + "BSHD on Q, K, V and O"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), {{"o", bshd, {}}}), + "O BSHD while Q, K and V are BHSD"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), onAllIoTensors(rowsPlus8)), + "row stride D+8 on Q, K, V and O"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), onAllIoTensors(rowsPlus2)), + "row stride D+2 on Q, K, V and O"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph({2, 2, 256, 192}, dims), + {{"q", {int64_t{2} * 256 * 194, int64_t{256} * 194, 194, 1}, {}}, + {"k", {int64_t{2} * 256 * 194, int64_t{256} * 194, 194, 1}, {}}, + {"v", rowsPlus2, {}}, + {"o", rowsPlus2, {}}}), + "hd192 with row stride D+2 on Q, K, V and O"}, + true); + // Unlike the backward kernels, the forward kernels addressed odd strides + // correctly on every tensor and axis, with bf16 on gfx942 and gfx950. The fp8 + // kernels were not run with them. + for(const char* name : K_IO_TENSOR_NAMES) + { + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), {{name, rowsPlus1, {}}}), + std::string(name) + " row stride D+1"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), {{name, oddHead, {}}}), + std::string(name) + " odd head stride"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), {{name, oddBatch, {}}}), + std::string(name) + " odd batch stride"}, + true); + } + // A single row of D elements fits any row stride of at least D. + applicabilityTests.emplace_back(GraphTest{createLayoutGraph({2, 2, 1, 128}, dims), "S_q=1"}, + true); + applicabilityTests.emplace_back(GraphTest{withLayouts(createLayoutGraph({2, 2, 1, 128}, dims), + {{"o", {256, 128, 256, 1}, {}}}), + "S_q=1 with O BSHD"}, + true); + applicabilityTests.emplace_back(GraphTest{createLayoutGraph(dims, {2, 2, 1, 128}), "S_kv=1"}, + true); + // A zero batch or head stride on an input (a broadcast) only re-reads one + // slice. + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), onTensors({"k", "v"}, zeroBatch)), + "K and V batch stride 0"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), onTensors({"k", "v"}, zeroHead)), + "K and V head stride 0"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims), {{"q", zeroBatch, {}}}), + "Q batch stride 0"}, + true); + // On a length-1 axis a zero stride does not alias, also on O. + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph({1, 2, 256, 128}, {1, 2, 256, 128}), + {{"o", zeroBatch, {}}}), + "O batch stride 0 with B=1"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph({2, 1, 256, 128}, {2, 1, 256, 128}), + {{"o", zeroHead, {}}}), + "O head stride 0 with H=1"}, + true); + + // Stats, which the kernels write at ((b * H_q + h) * head stride + s) * 4 bytes. + applicabilityTests.emplace_back(GraphTest{createLayoutGraph(dims, dims, true), "packed stats"}, + true); + applicabilityTests.emplace_back(GraphTest{withLayouts(createLayoutGraph(dims, dims, true), + {{"stats", {512, 256, 1}, {2, 2, 256}}}), + "rank-3 packed stats"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph(dims, dims, true), onAllIoTensors(bshd)), + "BSHD on Q, K, V and O with packed stats"}, + true); + applicabilityTests.emplace_back( + GraphTest{createLayoutGraph({2, 8, 256, 128}, {2, 2, 256, 128}, true), + "GQA with packed stats"}, + true); + applicabilityTests.emplace_back( + GraphTest{createLayoutGraph({2, 2, 256, 192}, dims, true), "hd192 with packed stats"}, + true); + applicabilityTests.emplace_back( + GraphTest{ + withLayouts(createLayoutGraph(dims, dims, true), {{"stats", {640, 320, 1, 1}, {}}}), + "stats head stride S_q + 64, batch stride H_q * (S_q + 64)"}, + true); + applicabilityTests.emplace_back( + GraphTest{createLayoutGraph({2, 1, 256, 128}, {2, 1, 256, 128}, true), + "packed stats with H_q=1"}, + true); + // No kernel uses the stride of a size-1 stats batch or sequence axis, nor the + // head stride when there is a single (batch, head) slice. + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph({1, 2, 256, 128}, {1, 2, 256, 128}, true), + {{"stats", {12345, 256, 1, 1}, {}}}), + "unused stats batch stride, B=1"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph({2, 2, 1, 128}, dims, true), + {{"stats", {2, 1, 5, 1}, {}}}), + "unused stats seq stride, S_q=1"}, + true); + applicabilityTests.emplace_back( + GraphTest{withLayouts(createLayoutGraph({1, 1, 256, 128}, {1, 1, 256, 128}, true), + {{"stats", {256, 7, 1, 1}, {}}}), + "unused stats head stride, B=1 H_q=1"}, + true); + applicabilityTests.emplace_back(GraphTest{withLayouts(createLayoutGraph(dims, dims, true), + {{"stats", {512, 256, 1, 7}, {}}}), + "arbitrary stride on the size-1 last stats dim"}, + true); + + for(const auto& [test, applicability] : applicabilityTests) + { + EXPECT_EQ(_planBuilder.isApplicable(_handle, test.graphWrapper()), applicability) + << test.message; + } +} + +// No kernel takes a head-dim stride, so a strided head dim on Q, K, V or O must +// decline. On gfx942 each of these gave a wrong O, and a stride-2 O left half of +// its elements unwritten. +TEST_F(TestSdpaFwdPlanBuilder, IsApplicableRejectsNonUnitHeadDimStride) +{ + SKIP_IF_NO_DEVICES(); + + const std::string deviceString + = hip_kernel_provider_common::getDeviceString(_handle.getStream()); + if(deviceString != "gfx942" && deviceString != "gfx950") + { + GTEST_SKIP(); + } + + const std::vector dims = {2, 2, 256, 128}; + // BHSD with head-dim stride 2 and row stride 2 * D. + const std::vector headDimStride2 + = {int64_t{2} * 256 * 256, int64_t{256} * 256, 256, 2}; + + std::vector declineTests; + declineTests.emplace_back( + withLayouts(createLayoutGraph({2, 2, 256, 192}, dims), + {{"q", {int64_t{2} * 256 * 384, int64_t{256} * 384, 384, 2}, {}}}), + "hd192 q head-dim stride 2"); + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), + {{"q", {int64_t{2} * 128 * 256, int64_t{128} * 256, 1, 256}, {}}}), + "q stored as [B, H, D, S]"); + // Strides of another rank are declined before any of them is read. + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), + {{"q", {int64_t{2} * 256 * 128, int64_t{256} * 128, 128, 1, 1}, {}}}), + "5 q strides for rank-4 q dims"); + for(const char* name : K_IO_TENSOR_NAMES) + { + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), {{name, headDimStride2, {}}}), + std::string(name) + " head-dim stride 2"); + } + + for(const auto& test : declineTests) + { + EXPECT_FALSE(_planBuilder.isApplicable(_handle, test.graphWrapper())) << test.message; + } +} + +// A row stride below the tensor's own head dim, also 0 and also on a length-1 +// sequence axis: each buffer ends S * row stride bytes into its (batch, head) +// slice, cutting off the end of the last row. On gfx942 the cases run gave a +// wrong O, an O of 0, or an O left NaN, except a K row stride of 1 with S_kv = 1: +// the one key takes all the softmax weight, so O was right, but by the address +// formula its score, and so the stats, are wrong. +TEST_F(TestSdpaFwdPlanBuilder, IsApplicableRejectsShortRows) +{ + SKIP_IF_NO_DEVICES(); + + const std::string deviceString + = hip_kernel_provider_common::getDeviceString(_handle.getStream()); + if(deviceString != "gfx942" && deviceString != "gfx950") + { + GTEST_SKIP(); + } + + const std::vector dims = {2, 2, 256, 128}; + // Rows overlap: 64-element row stride with D=128. + const std::vector overlappingRows + = {int64_t{2} * (int64_t{255} * 64 + 128), int64_t{255} * 64 + 128, 64, 1}; + // Every row of a (batch, head) slice at the same address. + const std::vector zeroRows = {256, 128, 0, 1}; + + std::vector declineTests; + for(const char* name : K_IO_TENSOR_NAMES) + { + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), {{name, overlappingRows, {}}}), + std::string(name) + " row stride D/2"); + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), {{name, zeroRows, {}}}), + std::string(name) + " row stride 0"); + } + // Q and O have S_q rows, K and V have S_kv rows. + for(const char* name : {"q", "o"}) + { + declineTests.emplace_back( + withLayouts(createLayoutGraph({2, 2, 1, 128}, dims), {{name, {256, 128, 1, 1}, {}}}), + std::string(name) + " row stride 1 with S_q=1"); + declineTests.emplace_back( + withLayouts(createLayoutGraph({2, 2, 2, 128}, dims), {{name, {384, 192, 64, 1}, {}}}), + std::string(name) + " row stride D/2 with S_q=2"); + } + for(const char* name : {"k", "v"}) + { + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, {2, 2, 1, 128}), {{name, {256, 128, 1, 1}, {}}}), + std::string(name) + " row stride 1 with S_kv=1"); + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, {2, 2, 2, 128}), {{name, {384, 192, 64, 1}, {}}}), + std::string(name) + " row stride D/2 with S_kv=2"); + } + // hd192x128: Q and K rows are 192 elements and V and O rows 128, so a row + // stride of 160 is too short for Q and K only. + const std::vector rows160 = {int64_t{2} * 256 * 160, int64_t{256} * 160, 160, 1}; + for(const char* name : {"q", "k"}) + { + declineTests.emplace_back( + withLayouts(createLayoutGraph({2, 2, 256, 192}, dims), {{name, rows160, {}}}), + std::string("hd192 ") + name + " row stride 160"); + } + + for(const auto& test : declineTests) + { + EXPECT_FALSE(_planBuilder.isApplicable(_handle, test.graphWrapper())) << test.message; + } +} + +// Each (batch, head) slice of O is written by its own workgroups, so a zero O +// batch or head stride on an axis longer than 1 makes slices overwrite each +// other. On gfx942 an O batch or head stride of 0 left half or two thirds of O +// wrong, and O changed between some runs with identical inputs. Only a zero +// stride is checked: nonzero strides that make O overlap itself, such as equal +// batch and head strides, are not. +TEST_F(TestSdpaFwdPlanBuilder, IsApplicableRejectsZeroStrideOutput) +{ + SKIP_IF_NO_DEVICES(); + + const std::string deviceString + = hip_kernel_provider_common::getDeviceString(_handle.getStream()); + if(deviceString != "gfx942" && deviceString != "gfx950") + { + GTEST_SKIP(); + } + + const std::vector dims = {2, 2, 256, 128}; + + std::vector declineTests; + // Only one axis aliases in each of these. + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), {{"o", {0, int64_t{256} * 128, 128, 1}, {}}}), + "O batch stride 0"); + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), {{"o", {int64_t{256} * 128, 0, 128, 1}, {}}}), + "O head stride 0"); + declineTests.emplace_back( + withLayouts(createLayoutGraph(dims, dims), {{"o", {0, 0, 128, 1}, {}}}), + "O batch and head stride 0"); + // The head stride may be 0 here, as H=1, but the batch stride may not. + declineTests.emplace_back(withLayouts(createLayoutGraph({2, 1, 256, 128}, {2, 1, 256, 128}), + {{"o", {0, 0, 128, 1}, {}}}), + "O batch and head stride 0, B=2 H=1"); + + for(const auto& test : declineTests) + { + EXPECT_FALSE(_planBuilder.isApplicable(_handle, test.graphWrapper())) << test.message; + } +} + +// Every forward kernel writes the LSE of (b, h, s) at +// ((b * H_q + h) * lse_Hs + s) * 4 bytes, with lse_Hs the stats head stride. On +// gfx942 each stats layout that this formula does not describe gave wrong or NaN +// stats in the cases run; by the formula, [H_q, B, S_q] storage also writes past +// the end of the stats buffer. A head stride below S_q makes (batch, head) slices +// overlap, so it declines too. +TEST_F(TestSdpaFwdPlanBuilder, IsApplicableRejectsUnaddressableStats) +{ + SKIP_IF_NO_DEVICES(); + + const std::string deviceString + = hip_kernel_provider_common::getDeviceString(_handle.getStream()); + if(deviceString != "gfx942" && deviceString != "gfx950") + { + GTEST_SKIP(); + } + + // Packed stats for these dims are [2, 2, 256, 1] with strides [512, 256, 1, 1]. + const std::vector dims = {2, 2, 256, 128}; + const auto statsGraph = [](const std::vector& qDims, + const std::vector& vDims, + const std::vector& statsStrides, + const std::vector& statsDims = {}) { + return withLayouts(createLayoutGraph(qDims, vDims, true), + {{"stats", statsStrides, statsDims}}); + }; + + std::vector declineTests; + // Only the seq stride is wrong in the next three. + declineTests.emplace_back(statsGraph(dims, dims, {512, 256, 2, 1}), "stats seq stride 2"); + declineTests.emplace_back(statsGraph(dims, dims, {512, 256, 0, 1}), + "stats broadcast along S_q"); + declineTests.emplace_back(statsGraph({1, 1, 256, 128}, {1, 1, 256, 128}, {512, 512, 2, 1}), + "stats seq stride 2, B=1 H_q=1"); + // Only the batch stride is wrong in the next four. + declineTests.emplace_back(statsGraph(dims, dims, {640, 256, 1, 1}), + "stats batch stride H_q * S_q + 128"); + declineTests.emplace_back(statsGraph(dims, dims, {256, 512, 1, 1}), + "stats stored as [H_q, B, S_q]"); + declineTests.emplace_back( + statsGraph({2, 8, 256, 128}, {2, 2, 256, 128}, {int64_t{8} * 256 + 128, 256, 1, 1}), + "GQA stats batch stride H_q * S_q + 128"); + declineTests.emplace_back(statsGraph({2, 2, 256, 192}, dims, {640, 256, 1, 1}), + "hd192 stats batch stride H_q * S_q + 128"); + // With H_q=1 the kernel still steps between batches by the head stride, so a + // batch stride that differs from it declines too. + declineTests.emplace_back(statsGraph({2, 1, 256, 128}, {2, 1, 256, 128}, {512, 256, 1, 1}), + "stats batch stride 2 * S_q, head stride S_q, H_q=1"); + // Only the head stride is wrong in the next four: (batch, head) slices overlap. + declineTests.emplace_back(statsGraph(dims, dims, {256, 128, 1, 1}), + "stats head stride S_q / 2"); + declineTests.emplace_back(statsGraph(dims, dims, {0, 0, 1, 1}), + "stats batch and head stride 0"); + declineTests.emplace_back(statsGraph({1, 2, 256, 128}, {1, 2, 256, 128}, {512, 0, 1, 1}), + "stats head stride 0, B=1"); + declineTests.emplace_back(statsGraph({2, 1, 256, 128}, {2, 1, 256, 128}, {128, 128, 1, 1}), + "stats batch and head stride S_q / 2, H_q=1"); + // More than one stride is wrong in the next three. With H_q=1, [S_q, 1, 1, 1] + // has both a batch stride other than the head stride and a head stride + // below S_q. + declineTests.emplace_back(statsGraph(dims, dims, {512, 1, 2, 1}), + "stats stored as [B, S_q, H_q]"); + declineTests.emplace_back(statsGraph(dims, dims, {512, 1, 2}, {2, 2, 256}), + "rank-3 stats stored as [B, S_q, H_q]"); + declineTests.emplace_back(statsGraph({2, 1, 256, 128}, {2, 1, 256, 128}, {256, 1, 1, 1}), + "stats strides [S_q, 1, 1, 1] with H_q=1"); + // Strides of another rank are declined before any of them is read. + declineTests.emplace_back(statsGraph(dims, dims, {512, 256, 1}), + "3 stats strides for rank-4 stats dims"); + + for(const auto& test : declineTests) + { + EXPECT_FALSE(_planBuilder.isApplicable(_handle, test.graphWrapper())) << test.message; + } +} + // ============================================================================= // buildPlan exception contract (IPlanBuilder::buildPlan) // =============================================================================