diff --git a/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaPlanUtils.hpp b/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaPlanUtils.hpp index 14d104312270..d02c07336ca9 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaPlanUtils.hpp +++ b/dnn-providers/hip-kernel-provider/src/engines/asm_sdpa_engine/plans/SdpaPlanUtils.hpp @@ -57,6 +57,12 @@ enum class MaskType : int // The diagonal band an SDPA (forward or backward) attribute set requests, with // the deprecated causal_mask / causal_mask_bottom_right booleans merged in. +// +// The bounds are in the kernels' FlashAttention convention (window_size_left / +// window_size_right), not the graph's: leftBound is the number of keys kept +// before the diagonal, so it is the graph's cuDNN-semantics left_bound minus 1 +// (see resolveDiagonalBand()). rightBound means the same in both conventions. +// -1 is unbounded. struct DiagonalBand { int64_t leftBound; @@ -100,7 +106,12 @@ inline DiagonalBand makeCoveringBoundsUnbound(DiagonalBand band, int64_t seqLenQ // alignment; causal_mask_bottom_right always resolves to BOTTOM_RIGHT, because // TOP_LEFT is the schema default and indistinguishable from an unset alignment. // -// Setting both booleans, or a bound below -1, throws +// The graph's left_bound follows cuDNN: left_bound = L keeps the L keys ending at +// the diagonal, diagonal included. The returned band holds the FlashAttention +// window_size_left the kernels take, L - 1 (see DiagonalBand). +// +// Setting both booleans, a bound below -1, or left_bound = 0 (which keeps no key at or +// before the diagonal, and which L - 1 would turn into "unbounded") throws // HipdnnPluginException(INVALID_VALUE). template DiagonalBand resolveDiagonalBand(const SdpaAttrsT& attrs) @@ -128,13 +139,22 @@ DiagonalBand resolveDiagonalBand(const SdpaAttrsT& attrs) "SDPA: left_bound and right_bound must be >= -1 (got left_bound=" + std::to_string(left) + ", right_bound=" + std::to_string(right) + ")"); } + if(left == 0) + { + throw hipdnn_plugin_sdk::HipdnnPluginException( + HIPDNN_PLUGIN_STATUS_INVALID_VALUE, + "SDPA: left_bound must be -1 (unbounded) or >= 1; left_bound = L keeps the L keys " + "ending at the diagonal, so 0 keeps no key at or before it"); + } + + const int64_t windowLeft = left > 0 ? left - 1 : -1; if(!causalDeprecated && !bottomRightDeprecated) { - return {left, right, attrs.diagonal_alignment()}; + return {windowLeft, right, attrs.diagonal_alignment()}; } - return {left, + return {windowLeft, right >= 0 ? right : 0, bottomRightDeprecated ? DiagonalAlignment::BOTTOM_RIGHT : attrs.diagonal_alignment()}; } @@ -163,8 +183,9 @@ inline MaskType getMaskType(int64_t leftBound, // Sliding-window mask coordinate transformation // ============================================================================= // -// Converts raw window sizes (left_bound, right_bound from the hipDNN graph) -// into the precomputed mask coordinates (mask_y, mask_x) that the AITER ASM +// Converts FlashAttention-convention window sizes (a DiagonalBand's bounds, not +// the graph's cuDNN-semantics left_bound) into the precomputed mask coordinates +// (mask_y, mask_x) that the AITER ASM // DQDKDV kernel expects in its argument struct. // // AITER reference: ck_tile_shim.h::compute_mask_coordinates() diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx1151WmmaAttentionNative.cpp b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx1151WmmaAttentionNative.cpp index e0b9f47111a2..fd4d8b068dec 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx1151WmmaAttentionNative.cpp +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx1151WmmaAttentionNative.cpp @@ -396,10 +396,15 @@ std::optional maskFor(const data_objects::SdpaAttributes& attrib = attributes.left_bound().has_value() ? attributes.left_bound().value() : UNBOUNDED; const int64_t right = attributes.right_bound().has_value() ? attributes.right_bound().value() : UNBOUNDED; - if(left < UNBOUNDED || right < UNBOUNDED || !fitsI32(left) || !fitsI32(right)) + // left_bound = 0 keeps no key at or before the diagonal; hipDNN rejects it. + if(left < UNBOUNDED || left == 0 || right < UNBOUNDED || !fitsI32(left) || !fitsI32(right)) { return std::nullopt; } + // The graph's left_bound follows cuDNN (L keys ending at the diagonal, diagonal + // included); the kernel's window_left follows FlashAttention (keys before the + // diagonal), so the same window is L - 1. + const int64_t windowLeft = left > 0 ? left - 1 : UNBOUNDED; MaskParameters result; // The deprecated causal flags override any bounds and alignment, exactly as @@ -413,11 +418,11 @@ std::optional maskFor(const data_objects::SdpaAttributes& attrib } result.bottomRight = static_cast(attributes.diagonal_alignment() == data_objects::DiagonalAlignment::BOTTOM_RIGHT); - result.left = static_cast(left); + result.left = static_cast(windowLeft); result.right = static_cast(right); - if(left != UNBOUNDED || right != UNBOUNDED) + if(windowLeft != UNBOUNDED || right != UNBOUNDED) { - result.kind = left == UNBOUNDED && right == 0 ? MaskKind::CAUSAL : MaskKind::WINDOW; + result.kind = windowLeft == UNBOUNDED && right == 0 ? MaskKind::CAUSAL : MaskKind::WINDOW; } return result; } diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/DeprecatedCausalMaskCases.hpp b/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/DeprecatedCausalMaskCases.hpp index 95a29b2308fc..fdbc8578589f 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/DeprecatedCausalMaskCases.hpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/DeprecatedCausalMaskCases.hpp @@ -224,6 +224,17 @@ inline const std::vector& deprecatedCausalMaskCases() std::nullopt}, DeprecatedCausalMaskCase{ "LeftBoundBelowMinusOne", false, false, -2, 0, ALIGN_TOP_LEFT, std::nullopt}, + // cuDNN semantics: left_bound = L keeps L keys, so 0 keeps no key at or before the + // diagonal and is rejected. + DeprecatedCausalMaskCase{"LeftBoundZero", false, false, 0, 0, ALIGN_TOP_LEFT, std::nullopt}, + // left_bound = 1 keeps only the diagonal: a band, not causal. + DeprecatedCausalMaskCase{"LeftBoundOneIsDiagonalOnly", + false, + false, + 1, + 0, + ALIGN_TOP_LEFT, + plan_utils::MaskType::SLIDING_WINDOW}, DeprecatedCausalMaskCase{"BoundsOnlyTopLeftCausal", false, false, diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaBwdPlanBuilder.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaBwdPlanBuilder.cpp index aabe8fb8d2d8..6ad9d2213fe0 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaBwdPlanBuilder.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/asm_sdpa_engine/TestSdpaBwdPlanBuilder.cpp @@ -1194,18 +1194,20 @@ TEST(TestSdpaBwdDeprecatedCausalMaskBand, CausalWithLeftBoundResolvesToCausalWin .front() ->attributesAs(); + // left_bound = 64 keeps 64 keys (cuDNN), which is FlashAttention's window_size_left 63. const auto band = plan_utils::resolveDiagonalBand(attrs); - EXPECT_EQ(band.leftBound, 64); + EXPECT_EQ(band.leftBound, 63); EXPECT_EQ(band.rightBound, 0); EXPECT_EQ(band.alignment, ALIGN_BOTTOM_RIGHT); constexpr int64_t SEQ_LEN_ABOVE_BOUND = 256; const auto windowed = plan_utils::makeCoveringBoundsUnbound(band, SEQ_LEN_ABOVE_BOUND, SEQ_LEN_ABOVE_BOUND); - EXPECT_EQ(windowed.leftBound, 64); + EXPECT_EQ(windowed.leftBound, 63); EXPECT_EQ(windowed.rightBound, 0); - constexpr int64_t SEQ_LEN_COVERED_BY_BOUND = 65; + // 64 keys cover every row of a 64-long sequence; at 65 the last row would drop key 0. + constexpr int64_t SEQ_LEN_COVERED_BY_BOUND = 64; const auto covered = plan_utils::makeCoveringBoundsUnbound( band, SEQ_LEN_COVERED_BY_BOUND, SEQ_LEN_COVERED_BY_BOUND); EXPECT_EQ(covered.leftBound, -1); 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 6979522f93ce..aab1e100930d 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 @@ -26,6 +26,7 @@ #include #include #include +#include namespace asm_sdpa_engine { @@ -875,6 +876,35 @@ INSTANTIATE_TEST_SUITE_P(, ::testing::ValuesIn(deprecatedCausalMaskCases()), deprecatedCausalMaskCaseName); +// The graph's left_bound follows cuDNN (L keys ending at the diagonal, diagonal +// included); the band the kernels take is FlashAttention's window_size_left, L - 1. +TEST(TestSdpaFwdResolveDiagonalBand, ConvertsCudnnLeftBoundToWindowLeft) +{ + // Literal expectations, so the test does not restate the conversion it checks. + for(const auto& [leftBound, expectedWindowLeft] : {std::pair{-1, -1}, + std::pair{1, 0}, + std::pair{2, 1}, + std::pair{64, 63}}) + { + auto builder = createSdpaFwdGraphWithMask( + /*causalMask=*/false, + /*causalMaskBottomRight=*/false, + leftBound, + /*rightBound=*/0, + hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment::TOP_LEFT); + const hipdnn_flatbuffers_sdk::flatbuffer_utilities::GraphWrapper graphWrapper( + builder.GetBufferPointer(), builder.GetSize()); + const auto& attrs + = graphWrapper.nodeWrappers() + .front() + ->attributesAs(); + + const auto band = plan_utils::resolveDiagonalBand(attrs); + EXPECT_EQ(band.leftBound, expectedWindowLeft) << "left_bound=" << leftBound; + EXPECT_EQ(band.rightBound, 0); + } +} + // ============================================================================= // plan_utils::makeCoveringBoundsUnbound // ============================================================================= diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/hip_flash2_engine/TestHipFlash2FwdPlanBuilder.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/hip_flash2_engine/TestHipFlash2FwdPlanBuilder.cpp index cd6110485da5..a8771038a94e 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/hip_flash2_engine/TestHipFlash2FwdPlanBuilder.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/hip_flash2_engine/TestHipFlash2FwdPlanBuilder.cpp @@ -334,8 +334,9 @@ INSTANTIATE_TEST_SUITE_P(, ::testing::ValuesIn(asm_sdpa_engine::deprecatedCausalMaskCases()), asm_sdpa_engine::deprecatedCausalMaskCaseName); -// The mask graph is 256x256, so a top-left left bound of 255 covers the whole -// sequence and is plain causal, while 254 is a genuine sliding window. +// The mask graph is 256x256. left_bound follows cuDNN (L keeps the L keys ending at +// the diagonal), so a top-left left bound of 256 covers the whole sequence and is +// plain causal, while 255 is a genuine sliding window (the last query drops key 0). bool extractCausalForTopLeftWindow(int64_t leftBound) { auto builder = asm_sdpa_engine::createSdpaFwdGraphWithMask( @@ -352,13 +353,13 @@ bool extractCausalForTopLeftWindow(int64_t leftBound) TEST(TestHipFlash2ExtractParamsWindow, CoveringLeftBoundIsCausal) { - EXPECT_TRUE(extractCausalForTopLeftWindow(255)); + EXPECT_TRUE(extractCausalForTopLeftWindow(256)); EXPECT_TRUE(extractCausalForTopLeftWindow(4294967298)); } TEST(TestHipFlash2ExtractParamsWindow, GenuineWindowIsNotCausal) { - EXPECT_FALSE(extractCausalForTopLeftWindow(254)); + EXPECT_FALSE(extractCausalForTopLeftWindow(255)); } } // namespace diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx1151WmmaAttentionMatchers.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx1151WmmaAttentionMatchers.cpp index 5fa8f40db613..939d5c70f1ce 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx1151WmmaAttentionMatchers.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx1151WmmaAttentionMatchers.cpp @@ -603,6 +603,32 @@ TEST(TestGfx1151WmmaAttentionMatchers, DeclinesOverrideShapeGraphs) EXPECT_FALSE(matchGraph(spec).has_value()); } +// The graph's left_bound follows cuDNN (L keys ending at the diagonal, diagonal +// included); the kernel's window_left follows FlashAttention (keys before the +// diagonal), so the same window binds as L - 1. left_bound = 0 is declined. +TEST(TestGfx1151WmmaAttentionMatchers, BindsCudnnLeftBoundAsFlashAttentionWindowLeft) +{ + using hipdnn_plugin_sdk::ingestor::tryGetBoundInt; + GraphSpec spec; + spec.leftBound = 7; + spec.rightBound = 3; + const auto bound = matchGraph(spec); + ASSERT_TRUE(bound.has_value()); + EXPECT_EQ(tryGetBoundInt(*bound, "gfx1151_wmma_attention.mask"), 2); // window + EXPECT_EQ(tryGetBoundInt(*bound, "gfx1151_wmma_attention.window_left"), 6); + EXPECT_EQ(tryGetBoundInt(*bound, "gfx1151_wmma_attention.window_right"), 3); + + spec.leftBound = 1; + spec.rightBound = 0; + const auto diagonalOnly = matchGraph(spec); + ASSERT_TRUE(diagonalOnly.has_value()); + EXPECT_EQ(tryGetBoundInt(*diagonalOnly, "gfx1151_wmma_attention.mask"), 2); // window + EXPECT_EQ(tryGetBoundInt(*diagonalOnly, "gfx1151_wmma_attention.window_left"), 0); + + spec.leftBound = 0; + EXPECT_FALSE(matchGraph(spec).has_value()); +} + // The deprecated causal flags override the bounds and alignment, as in the // hipDNN reference; otherwise a left bound would drop the causal limit. TEST(TestGfx1151WmmaAttentionMatchers, DeprecatedCausalFlagsOverrideBounds) @@ -942,7 +968,7 @@ TEST(TestGfx1151WmmaAttentionMatchers, DecodeObjectServesTheDispatcherDecodeRule window.kvHeads = 8; EXPECT_TRUE(matchesKernel(window, decodeKernel(128), gfx12)); window.bottomRight = true; - window.leftBound = 1023; + window.leftBound = 1024; // 1024 keys including the diagonal: window_left 1023 window.rightBound = 0; EXPECT_FALSE(matchesKernel(window, decodeKernel(128), gfx12)); } @@ -1004,7 +1030,7 @@ TEST(TestGfx1151WmmaAttentionMatchers, DecodeWorkspaceMirrorsTheDispatcherPlan) // A bottom-right 1024-key window: 2 splits (still merged by the merge kernel). GraphSpec window = decodeGraph(); window.bottomRight = true; - window.leftBound = 1023; + window.leftBound = 1024; // 1024 keys including the diagonal: window_left 1023 window.rightBound = 0; EXPECT_EQ(workspaceOf(window, decodeKernel(128), gfx1201), 33024U); // Batch 4: 2 splits. diff --git a/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaFwd.cpp b/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaFwd.cpp index d928a6668d81..bc0fe5a1ab61 100644 --- a/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaFwd.cpp +++ b/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaFwd.cpp @@ -126,8 +126,9 @@ extern "C" __global__ void sdpaFwdRef(SdpaFwdArgs args) s += toAccum(mask[maskOffset]); } - // (b) Sliding-window mask (applied after the additive mask). - // Asymmetric: +1 on the right bound, none on the left bound. + // (b) Sliding-window mask (applied after the additive mask), cuDNN semantics: + // keeps skv in [sq + windowOffset - leftBound + 1, sq + windowOffset + rightBound], + // so leftBound = L keeps the L keys ending at the diagonal. if(args.rightBound >= 0) { long long startKv = sq + 1 + windowOffset + args.rightBound; @@ -142,7 +143,7 @@ extern "C" __global__ void sdpaFwdRef(SdpaFwdArgs args) } if(args.leftBound >= 0) { - if(skv < sq + windowOffset - args.leftBound) + if(skv <= sq + windowOffset - args.leftBound) { s = negInf; } diff --git a/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaRaggedFwd.cpp b/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaRaggedFwd.cpp index 5918a6cac705..ea383444d670 100644 --- a/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaRaggedFwd.cpp +++ b/dnn-providers/integration-tests/gpu-ref/kernels/sdpa/GpuRefSdpaRaggedFwd.cpp @@ -182,7 +182,8 @@ extern "C" __global__ void sdpaRaggedFwdRef(SdpaRaggedFwdArgs args) // Fold in the fp8 Q/K descale. COMPUTE_TYPE s = dot * descaleQK * static_cast(args.scale); - // Sliding-window mask, aligned per batch. Only the right bound gets the +1. + // Sliding-window mask, aligned per batch, cuDNN semantics: keeps skv in + // [sq + windowOffset - leftBound + 1, sq + windowOffset + rightBound]. // No additive bias: it is gated off on the ASM v3 path. if(args.rightBound >= 0) { @@ -198,7 +199,7 @@ extern "C" __global__ void sdpaRaggedFwdRef(SdpaRaggedFwdArgs args) } if(args.leftBound >= 0) { - if(skv < sq + windowOffset - args.leftBound) + if(skv <= sq + windowOffset - args.leftBound) { s = negInf; } diff --git a/dnn-providers/integration-tests/reference-data-scripts/README.md b/dnn-providers/integration-tests/reference-data-scripts/README.md index 7984cc5c86e5..e18b32757c8d 100644 --- a/dnn-providers/integration-tests/reference-data-scripts/README.md +++ b/dnn-providers/integration-tests/reference-data-scripts/README.md @@ -124,3 +124,18 @@ and for template-sweep bundles: - `golden.path` validation and per-case tensor file/NaN/size checks (input and output tensors) - per-case metadata field validation - a directory with only one of `graph.template.json`/`sweep.json` + +## Sliding-window bounds in the generators + +A left window bound has two common conventions, and each generator flag is named after the +one it uses: + +| Convention | Name | `L = 3` keeps, for the query on the diagonal | +|---|---|---| +| hipDNN graph, cuDNN | `left_bound` (`--left-bound`) | the 3 keys ending at the diagonal, diagonal included | +| FlashAttention, AITER | `window_size_left` (`--window-left`) | the 3 keys before the diagonal, plus the diagonal (4 keys) | + +`generate_sdpa_fwd_golden.py --left-bound L` writes `left_bound = L` to the graph. +`--window-left W` writes `left_bound = W + 1`; it exists for `--validate`, which passes the +window to AITER. The two flags are mutually exclusive. `left_bound` must be positive or -1 +(unbounded). diff --git a/dnn-providers/integration-tests/reference-data-scripts/generate_sdpa_fwd_golden.py b/dnn-providers/integration-tests/reference-data-scripts/generate_sdpa_fwd_golden.py index 5b5d72551979..ea3b0e96380b 100644 --- a/dnn-providers/integration-tests/reference-data-scripts/generate_sdpa_fwd_golden.py +++ b/dnn-providers/integration-tests/reference-data-scripts/generate_sdpa_fwd_golden.py @@ -13,7 +13,8 @@ - dtype: bf16, fp16, fp8 (FP8_E4M3 inputs, BF16 output, per-tensor descale) - head dim: hd128, hd192 (via --q-dims / --v-dims) - causal mask: none, top_left, bottom_right - - sliding window: --window-left / --window-right + - sliding window: --left-bound (cuDNN / graph convention) or --window-left + (FlashAttention convention), and --window-right - GROUP mode: --seq-lens-q / --seq-lens-kv (actual per-batch valid lengths) - ragged: --ragged-offsets (RFC 0014 packed BSHD; requires --layout bshd) - LSE output: --stats (adds tensor uid=4) @@ -186,13 +187,36 @@ def physical_seqlens(valid_lens, s_max): return [min(s_max, v + RAGGED_PAD_ROWS) for v in valid_lens] +def left_bound_from_window_left(window_left): + """Map a FlashAttention-style window_size_left to hipDNN's left_bound. + + hipDNN follows cuDNN: left_bound = L keeps the L keys ending at the diagonal, + diagonal included. FlashAttention's window_size_left = W keeps W + 1 keys, so + the same window is left_bound = W + 1. + """ + return window_left + 1 + + +def window_left_from_left_bound(left_bound): + """Map hipDNN's (cuDNN) left_bound to FlashAttention's window_size_left. + + The inverse of left_bound_from_window_left: L keys including the diagonal is + L - 1 keys before it. -1 (unbounded) stays -1. + """ + return -1 if left_bound == -1 else left_bound - 1 + + def build_mask(s_q, s_kv, causal, window_left, window_right): """Construct a boolean keep-mask [S_q, S_kv] for the given causal/window config. Returns None when no masking applies (the dense attention case), so callers can take the fast PyTorch path. The mask is True where attention is allowed. - The bound math matches the CPU reference executor's leftBound/rightBound + - topLeftAlignment semantics so golden output is reproducible on the CPU side. + + window_left / window_right follow FlashAttention's window_size_left/right: + window_left = W keeps the W keys before the diagonal plus the diagonal itself + (W + 1 keys). The graph's cuDNN-semantics left_bound for that window is W + 1 + (see left_bound_from_window_left), so the CPU reference executor reproduces + this mask. causal: - "none": no causal constraint @@ -402,7 +426,7 @@ def build_graph_json( left_bound = -1 right_bound = 0 if window_left >= 0: - left_bound = window_left + left_bound = left_bound_from_window_left(window_left) if window_right >= 0: right_bound = window_right @@ -1076,11 +1100,25 @@ def main(): help="Causal mask type. bottom_right maps to AITER mask=2; top_left is " "reference-only (no AITER kernel) (default: none)", ) - parser.add_argument( + # Two names for the same bound, one per convention, so a value is never read in the + # wrong one: --left-bound is the graph's (cuDNN) left_bound, --window-left is + # FlashAttention's window_size_left (what --validate passes to AITER). + left = parser.add_mutually_exclusive_group() + left.add_argument( + "--left-bound", + type=int, + default=None, + help="Sliding window left bound in the graph's (cuDNN) convention: L keeps the L " + "keys ending at the diagonal, diagonal included. Written to the graph unchanged. " + "Positive or -1 (unbounded); 0 is rejected, as hipDNN rejects it", + ) + left.add_argument( "--window-left", type=int, default=-1, - help="Sliding window left bound, -1 = unbounded (default: -1)", + help="Sliding window left size in FlashAttention's window_size_left convention " + "(W keeps the W keys before the diagonal plus the diagonal; written to the graph " + "as left_bound = W + 1), -1 = unbounded (default: -1)", ) parser.add_argument( "--window-right", @@ -1146,6 +1184,12 @@ def main(): ) args = parser.parse_args() + window_left = args.window_left + if args.left_bound is not None: + if args.left_bound == 0 or args.left_bound < -1: + parser.error("--left-bound must be positive or -1 (unbounded)") + window_left = window_left_from_left_bound(args.left_bound) + if args.install_aiter: install_aiter() @@ -1155,7 +1199,7 @@ def main(): v_dims=args.v_dims, dtype=args.dtype, causal=args.causal, - window_left=args.window_left, + window_left=window_left, window_right=args.window_right, seq_lens_q=args.seq_lens_q, seq_lens_kv=args.seq_lens_kv, diff --git a/dnn-providers/integration-tests/reference-data-scripts/tests/test_generate_sdpa_fwd_golden.py b/dnn-providers/integration-tests/reference-data-scripts/tests/test_generate_sdpa_fwd_golden.py index 5a5ee41c440e..b0f1e23a0c55 100644 --- a/dnn-providers/integration-tests/reference-data-scripts/tests/test_generate_sdpa_fwd_golden.py +++ b/dnn-providers/integration-tests/reference-data-scripts/tests/test_generate_sdpa_fwd_golden.py @@ -62,6 +62,108 @@ def test_explicit_attn_scale_is_kept(self): self.assertEqual(graph["nodes"][0]["attributes"]["attn_scale_value"], 0.25) self.assertEqual(meta["config"]["scale"], 0.25) + def test_window_left_is_written_as_cudnn_left_bound(self): + # --window-left is FlashAttention's window_size_left (W keeps W + 1 keys); + # the graph's left_bound follows cuDNN (L keeps L keys), so it is W + 1. + graph, meta = self.generate("--causal", "top_left", "--window-left", "2") + attributes = graph["nodes"][0]["attributes"] + self.assertEqual(attributes["left_bound"], 3) + self.assertEqual(attributes["right_bound"], 0) + self.assertEqual(meta["config"]["window_left"], 2) + + def test_left_bound_is_written_unchanged(self): + # --left-bound is the graph's (cuDNN) left_bound: the same window as --window-left 2. + graph, meta = self.generate("--causal", "top_left", "--left-bound", "3") + self.assertEqual(graph["nodes"][0]["attributes"]["left_bound"], 3) + self.assertEqual(meta["config"]["window_left"], 2) + window_graph, _ = self.generate("--causal", "top_left", "--window-left", "2") + self.assertEqual(graph["nodes"], window_graph["nodes"]) + + def test_left_bound_one_is_diagonal_only(self): + graph, meta = self.generate("--left-bound", "1", "--window-right", "0") + self.assertEqual(graph["nodes"][0]["attributes"]["left_bound"], 1) + self.assertEqual(meta["config"]["window_left"], 0) + + def test_invalid_left_bound_flags_are_rejected(self): + for extra in ( + ("--left-bound", "0"), + ("--left-bound", "-2"), + ("--left-bound", "3", "--window-left", "2"), + ): + with self.subTest(extra=extra), self.assertRaises( + subprocess.CalledProcessError + ): + self.generate(*extra) + + +def cudnn_keep_mask(torch, s_q, s_kv, diag_offset, left_bound, right_bound): + """cuDNN's sliding-window rule, as cudnn-frontend's own reference states it + (test/python/sdpa/fp16_ref.py): with rel = col - row, a key is masked when + rel > diag_offset + right_bound or rel <= diag_offset - left_bound. It shares no + code or convention with hipDNN's references, so it pins the contract from outside. + """ + rows = torch.arange(s_q).view(s_q, 1) + cols = torch.arange(s_kv).view(1, s_kv) + rel = cols - rows + masked = torch.zeros(s_q, s_kv, dtype=torch.bool) + if right_bound is not None: + masked |= rel > diag_offset + right_bound + if left_bound is not None: + masked |= rel <= diag_offset - left_bound + return ~masked + + +@unittest.skipIf(TORCH_MISSING, "the generator requires PyTorch") +class TestGenerateSdpaFwdGoldenWindowMatchesCudnn(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.gen = load_generator_module() + + def assert_matches_cudnn(self, s_q, s_kv, causal, window_left, window_right): + torch = self.gen.torch + keep = self.gen.build_mask(s_q, s_kv, causal, window_left, window_right) + diag_offset = s_kv - s_q if causal == "bottom_right" else 0 + if causal != "none": + right_bound = 0 + else: + right_bound = window_right if window_right >= 0 else None + expected = cudnn_keep_mask( + torch, + s_q, + s_kv, + diag_offset, + self.gen.left_bound_from_window_left(window_left), + right_bound, + ) + self.assertTrue( + torch.equal(keep, expected), + f"S_q={s_q} S_kv={s_kv} causal={causal} window=({window_left}," + f"{window_right})\nbuild_mask:\n{keep.int()}\ncuDNN:\n{expected.int()}", + ) + + def test_left_bound_two_keeps_two_keys_including_the_diagonal(self): + # cuDNN: left_bound = 2, TOP_LEFT, query 5 keeps keys {4, 5}. + torch = self.gen.torch + keep = cudnn_keep_mask(torch, 6, 6, 0, 2, 0) + self.assertEqual(keep[5].nonzero().flatten().tolist(), [4, 5]) + self.assertEqual(self.gen.left_bound_from_window_left(1), 2) + self.assertEqual(self.gen.window_left_from_left_bound(2), 1) + self.assertEqual(self.gen.window_left_from_left_bound(-1), -1) + + def test_causal_windows_match_cudnn(self): + for s_q, s_kv in ((6, 6), (4, 7), (7, 4)): + for causal in ("top_left", "bottom_right"): + for window_left in (0, 1, 2, 5): + self.assert_matches_cudnn(s_q, s_kv, causal, window_left, -1) + + def test_non_causal_windows_match_cudnn(self): + for s_q, s_kv in ((6, 6), (4, 7)): + for window_left in (0, 1, 2, 5): + for window_right in (-1, 0, 2): + self.assert_matches_cudnn( + s_q, s_kv, "none", window_left, window_right + ) + @unittest.skipIf(TORCH_MISSING, "the generator requires PyTorch") class TestGenerateSdpaFwdGoldenRaggedHelpers(unittest.TestCase): diff --git a/dnn-providers/integration-tests/src/harness/gpu-graph-executor/detail/GpuSdpaRaggedFwdPlan.hpp b/dnn-providers/integration-tests/src/harness/gpu-graph-executor/detail/GpuSdpaRaggedFwdPlan.hpp index 3e52c376829a..ba19009976f3 100644 --- a/dnn-providers/integration-tests/src/harness/gpu-graph-executor/detail/GpuSdpaRaggedFwdPlan.hpp +++ b/dnn-providers/integration-tests/src/harness/gpu-graph-executor/detail/GpuSdpaRaggedFwdPlan.hpp @@ -477,6 +477,12 @@ class GpuSdpaRaggedFwdPlanBuilder : public IGpuGraphNodePlanBuilder "GpuSdpaRaggedFwdPlan: left_bound and right_bound must be >= -1 (got left_bound=" + std::to_string(leftBound) + ", right_bound=" + std::to_string(rightBound) + ")"); } + // left_bound follows cuDNN: L keeps the L keys ending at the diagonal; 0 is rejected. + if(leftBound == 0) + { + throw std::invalid_argument( + "GpuSdpaRaggedFwdPlan: left_bound must be -1 (unbounded) or >= 1 (got 0)"); + } bool isTopLeft = nodeAttributes->diagonal_alignment() == hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment::TOP_LEFT; diff --git a/dnn-providers/integration-tests/tests/TestGpuSdpaRaggedFwdPlan.cpp b/dnn-providers/integration-tests/tests/TestGpuSdpaRaggedFwdPlan.cpp index 3e5ed51cf0a0..ee1304c4360d 100644 --- a/dnn-providers/integration-tests/tests/TestGpuSdpaRaggedFwdPlan.cpp +++ b/dnn-providers/integration-tests/tests/TestGpuSdpaRaggedFwdPlan.cpp @@ -1307,6 +1307,15 @@ TEST(TestGpuSdpaRaggedFwdPlanBuilder, BuildRejectsInvalidMaskAttributes) EXPECT_THROW(bf16Builder.buildNodePlan(badBoundWrap, badBoundWrap.getNode(0)), std::invalid_argument); + // cuDNN's left_bound counts the diagonal, so 0 keeps no key at or before it. + SdpaAttributesT zeroLeftBound; + zeroLeftBound.left_bound = 0; + auto zeroLeftGraph = makeRaggedGraph(zeroLeftBound); + auto zeroLeftWrap = hipdnn_flatbuffers_sdk::flatbuffer_utilities::GraphWrapper( + zeroLeftGraph.GetBufferPointer(), zeroLeftGraph.GetSize()); + EXPECT_THROW(bf16Builder.buildNodePlan(zeroLeftWrap, zeroLeftWrap.getNode(0)), + std::invalid_argument); + SdpaAttributesT badRightBound; badRightBound.right_bound = -2; auto badRightGraph = makeRaggedGraph(badRightBound); diff --git a/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaAttributes.hpp b/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaAttributes.hpp index 9cc99c9cb683..987668aefcd6 100644 --- a/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaAttributes.hpp +++ b/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaAttributes.hpp @@ -667,6 +667,22 @@ class SdpaAttributes : public Attributes attn_scale_value = value; return *this; } + /** + * @brief Set the left bound of the diagonal attention band (sliding window). + * + * Follows cuDNN: a left bound of L keeps the L keys ending at the diagonal, diagonal + * included, so query i attends to keys j with j > i + offset - L (offset is 0 for + * TOP_LEFT and seq_kv - seq_q for BOTTOM_RIGHT). Must be >= 1, or -1 for unbounded. + * A FlashAttention-style window_size_left = W is the same window as a left bound of W + 1. + * + * Two values differ from cuDNN. -1 means unbounded here, where cuDNN leaves the bound + * unset. 0 is rejected in both directions, where cuDNN 9.10 and later accepts it in the + * forward pass: a provider whose kernels take FlashAttention's window would pass L - 1, + * turning 0 into -1, which those kernels read as unbounded. + * + * @param value The left bound L. + * @return Reference to this for chaining. + */ // NOLINTNEXTLINE(readability-identifier-naming) SdpaAttributes& set_diagonal_band_left_bound(int64_t value) { @@ -746,7 +762,14 @@ class SdpaAttributes : public Attributes { return set_generate_stats(!value); } - // cuDNN set_sliding_window_length(n) forwards to the left diagonal-band bound. + /** + * @brief cuDNN-compatible alias of set_diagonal_band_left_bound(). + * + * A window length of n keeps the n keys ending at the diagonal, as in cuDNN. + * + * @param value The window length n (>= 1). + * @return Reference to this for chaining. + */ // NOLINTNEXTLINE(readability-identifier-naming) SdpaAttributes& set_sliding_window_length(int value) { diff --git a/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaBackwardAttributes.hpp b/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaBackwardAttributes.hpp index a51309970911..d1443b3c2cb2 100644 --- a/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaBackwardAttributes.hpp +++ b/projects/hipdnn/frontend/include/hipdnn_frontend/attributes/SdpaBackwardAttributes.hpp @@ -495,6 +495,22 @@ class SdpaBackwardAttributes : public Attributes setInput(InputNames::Offset, offset); return *this; } + /** + * @brief Set the left bound of the diagonal attention band (sliding window). + * + * Follows cuDNN: a left bound of L keeps the L keys ending at the diagonal, diagonal + * included, so query i attends to keys j with j > i + offset - L (offset is 0 for + * TOP_LEFT and seq_kv - seq_q for BOTTOM_RIGHT). Must be >= 1, or -1 for unbounded. + * A FlashAttention-style window_size_left = W is the same window as a left bound of W + 1. + * + * Two values differ from cuDNN. -1 means unbounded here, where cuDNN leaves the bound + * unset. 0 is rejected in both directions, where cuDNN 9.10 and later accepts it in the + * forward pass: a provider whose kernels take FlashAttention's window would pass L - 1, + * turning 0 into -1, which those kernels read as unbounded. + * + * @param value The left bound L. + * @return Reference to this for chaining. + */ // NOLINTNEXTLINE(readability-identifier-naming) SdpaBackwardAttributes& set_diagonal_band_left_bound(int64_t value) { @@ -544,7 +560,14 @@ class SdpaBackwardAttributes : public Attributes // than a straight rename of a native setter (overload merge, one-to-many // split, or a capability hipDNN lacks). - // cuDNN set_sliding_window_length(n) forwards to the left diagonal-band bound. + /** + * @brief cuDNN-compatible alias of set_diagonal_band_left_bound(). + * + * A window length of n keeps the n keys ending at the diagonal, as in cuDNN. + * + * @param value The window length n (>= 1). + * @return Reference to this for chaining. + */ // NOLINTNEXTLINE(readability-identifier-naming) SdpaBackwardAttributes& set_sliding_window_length(int value) { diff --git a/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaBwdNode.hpp b/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaBwdNode.hpp index 4f961603ff81..a8e2d817a4f4 100644 --- a/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaBwdNode.hpp +++ b/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaBwdNode.hpp @@ -357,6 +357,9 @@ class SdpaBwdNode : public BaseNode HIPDNN_CHECK_ERROR(detail::validateScalarParameter(scale, "SCALE tensor")); } + // Rule 11: left_bound follows cuDNN and must be -1 or >= 1 + HIPDNN_CHECK_ERROR(detail::validateSdpaLeftBound(attributes.left_bound, "SdpaBwdNode")); + return {}; } diff --git a/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaFwdNode.hpp b/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaFwdNode.hpp index a63976ec2f28..8c869b653c1d 100644 --- a/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaFwdNode.hpp +++ b/projects/hipdnn/frontend/include/hipdnn_frontend/node/SdpaFwdNode.hpp @@ -165,6 +165,9 @@ class SdpaFwdNode : public BaseNode HIPDNN_CHECK_ERROR(detail::validateScalarParameter(scale, "SCALE tensor")); } + // Rule 6c: left_bound follows cuDNN and must be -1 or >= 1 + HIPDNN_CHECK_ERROR(detail::validateSdpaLeftBound(attributes.left_bound, "SdpaFwdNode")); + // Rule 7: If O has dims already set, validate them const auto& oDims = o->get_dim(); if(!oDims.empty()) diff --git a/projects/hipdnn/frontend/include/hipdnn_frontend/node/detail/Utilities.hpp b/projects/hipdnn/frontend/include/hipdnn_frontend/node/detail/Utilities.hpp index 203e1c59a016..f6f850cefe84 100644 --- a/projects/hipdnn/frontend/include/hipdnn_frontend/node/detail/Utilities.hpp +++ b/projects/hipdnn/frontend/include/hipdnn_frontend/node/detail/Utilities.hpp @@ -7,6 +7,8 @@ #include #include #include +#include +#include #include namespace hipdnn_frontend::detail @@ -469,6 +471,24 @@ inline Error validateNormStatsShapeIfSet(const std::shared_ptr& leftBound, + const std::string& nodeName) +{ + HIPDNN_RETURN_IF_TRUE(leftBound.has_value() + && (leftBound.value() == 0 || leftBound.value() < -1), + ErrorCode::INVALID_VALUE, + nodeName + ": left_bound must be -1 (unbounded) or >= 1, got " + + std::to_string(leftBound.value_or(0)) + + ". left_bound = L keeps the L keys ending at the diagonal, " + "diagonal included"); + return {}; +} + // Validates scalar parameter tensor is properly configured // Uses param's name if set, otherwise uses fallbackName for error messages // Used for required scalar parameters (e.g., epsilon) that must have dimensions set diff --git a/projects/hipdnn/frontend/tests/TestSdpaBwdNode.cpp b/projects/hipdnn/frontend/tests/TestSdpaBwdNode.cpp index 39c3f648041a..f6e2f516de22 100644 --- a/projects/hipdnn/frontend/tests/TestSdpaBwdNode.cpp +++ b/projects/hipdnn/frontend/tests/TestSdpaBwdNode.cpp @@ -962,3 +962,29 @@ TEST(TestSdpaBwdNode, PreValidateFailsEachUnsupportedSetter) [](SdpaBackwardAttributes& a) { a.set_dsink_token(std::make_shared()); }, "sink-token gradient"); } + +// left_bound follows cuDNN: L keeps the L keys ending at the diagonal; 0 is rejected and +// -1 is unbounded. +TEST(TestSdpaBwdNode, PreValidateChecksLeftBound) +{ + const auto validateWithLeftBound = [](int64_t leftBound) { + auto attrs = makeMinimalAttrs( + makeTensor4D(2, 8, 16, 64), makeTensor4D(2, 8, 32, 64), makeTensor4D(2, 8, 32, 64)); + attrs.set_diagonal_band_left_bound(leftBound); + const GraphAttributes graphAttrs; + const SdpaBwdNode node(std::move(attrs), graphAttrs); + return node.pre_validate_node(); + }; + + for(const int64_t valid : {int64_t{-1}, int64_t{1}, int64_t{64}}) + { + const auto err = validateWithLeftBound(valid); + EXPECT_EQ(err.code, error_code_t::OK) << "left_bound=" << valid << ": " << err.err_msg; + } + for(const int64_t invalid : {int64_t{0}, int64_t{-2}}) + { + const auto err = validateWithLeftBound(invalid); + EXPECT_EQ(err.code, error_code_t::INVALID_VALUE) << "left_bound=" << invalid; + EXPECT_NE(err.err_msg.find("left_bound"), std::string::npos) << err.err_msg; + } +} diff --git a/projects/hipdnn/frontend/tests/TestSdpaFwdNode.cpp b/projects/hipdnn/frontend/tests/TestSdpaFwdNode.cpp index bc9a3a64ddd9..3579d1831897 100644 --- a/projects/hipdnn/frontend/tests/TestSdpaFwdNode.cpp +++ b/projects/hipdnn/frontend/tests/TestSdpaFwdNode.cpp @@ -574,3 +574,29 @@ TEST(TestSdpaFwdNode, PreValidateDrainRunsBeforeShapeChecks) EXPECT_NE(err.err_msg.find("score modifier"), std::string::npos) << err.err_msg; EXPECT_EQ(err.err_msg.find("rank"), std::string::npos) << err.err_msg; } + +// left_bound follows cuDNN: L keeps the L keys ending at the diagonal; 0 is rejected and +// -1 is unbounded. +TEST(TestSdpaFwdNode, PreValidateChecksLeftBound) +{ + const auto validateWithLeftBound = [](int64_t leftBound) { + auto attrs = makeMinimalAttrs( + makeTensor4D(2, 8, 16, 64), makeTensor4D(2, 8, 32, 64), makeTensor4D(2, 8, 32, 64)); + attrs.set_diagonal_band_left_bound(leftBound); + const GraphAttributes graphAttrs; + const SdpaFwdNode node(std::move(attrs), graphAttrs); + return node.pre_validate_node(); + }; + + for(const int64_t valid : {int64_t{-1}, int64_t{1}, int64_t{64}}) + { + const auto err = validateWithLeftBound(valid); + EXPECT_EQ(err.code, error_code_t::OK) << "left_bound=" << valid << ": " << err.err_msg; + } + for(const int64_t invalid : {int64_t{0}, int64_t{-2}}) + { + const auto err = validateWithLeftBound(invalid); + EXPECT_EQ(err.code, error_code_t::INVALID_VALUE) << "left_bound=" << invalid; + EXPECT_NE(err.err_msg.find("left_bound"), std::string::npos) << err.err_msg; + } +} diff --git a/projects/hipdnn/plugin_sdk/tests/heuristics/TestEngineFeatures.cpp b/projects/hipdnn/plugin_sdk/tests/heuristics/TestEngineFeatures.cpp index a47c2a513d97..e90de1c1c85b 100644 --- a/projects/hipdnn/plugin_sdk/tests/heuristics/TestEngineFeatures.cpp +++ b/projects/hipdnn/plugin_sdk/tests/heuristics/TestEngineFeatures.cpp @@ -1244,7 +1244,7 @@ TEST(TestEngineFeatures, WorkModelMatchesEveryGoldenCount) {"dropout_mask", [](SdpaAttributesT& op) { op.dropout_mask_tensor_uid = 5; }}, {"dropout_scale", [](SdpaAttributesT& op) { op.dropout_scale_tensor_uid = 5; }}, {"dropout_probability", [](SdpaAttributesT& op) { op.dropout_probability = 0.1F; }}, - {"left_bound", [](SdpaAttributesT& op) { op.left_bound = 0; }}, + {"left_bound", [](SdpaAttributesT& op) { op.left_bound = 1; }}, {"right_bound", [](SdpaAttributesT& op) { op.right_bound = 1; }}, }; for(const auto& [name, refuse] : refusals) diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpa.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpa.hpp index d64202cd325b..35424b52ab9b 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpa.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpa.hpp @@ -42,8 +42,8 @@ class CpuFpReferenceSdpa /// @param attnScaleValue Optional scale factor; 1.0 (no scaling) when unset, as in cuDNN /// @param attnMask Optional additive attention mask with dims right-aligned /// to [B, H, Sq, Skv] (rank 1–4), with broadcasting on size-1 dims - /// @param leftBound Number of sequence elements to the left that are unmasked - /// -1 is no mask to the left + /// @param leftBound Number of unmasked sequence elements ending at the diagonal, + /// diagonal included (cuDNN semantics). -1 is no mask to the left /// @param rightBound Number of sequence elements to the right that are unmasked /// -1 is no mask to the right /// @param topLeftAlignment If true, diagonal is measured from top left @@ -428,7 +428,8 @@ class CpuFpReferenceSdpa /// When provided, enables efficient softmax recomputation. /// When nullptr, recomputes softmax from scratch. /// @param attnMask Optional additive attention mask (same as forward) - /// @param leftBound Number of KV positions to the left that are unmasked (-1 = no mask) + /// @param leftBound Number of unmasked KV positions ending at the diagonal, diagonal + /// included (cuDNN semantics; -1 = no mask) /// @param rightBound Number of KV positions to the right that are unmasked (-1 = no mask) /// @param topLeftAlignment If true, diagonal at skv==sq; if false, at skv==sq+(seqKv-seqQ) /// @@ -880,6 +881,8 @@ class CpuFpReferenceSdpa /// leftBound = -1 means no left mask; rightBound = -1 means no right mask. /// topLeftAlignment: if true, diagonal at skv==sq; if false, diagonal at /// skv==sq+(seqKv-seqQ). + /// Follows cuDNN: leftBound = L keeps the L keys ending at the diagonal (diagonal + /// included), and rightBound = R keeps R keys past it. static bool isMasked(int64_t sq, int64_t skv, int64_t seqQ, @@ -893,7 +896,7 @@ class CpuFpReferenceSdpa { return true; } - if(leftBound >= 0 && skv < sq + offset - leftBound) + if(leftBound >= 0 && skv <= sq + offset - leftBound) { return true; } diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpaRagged.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpaRagged.hpp index 80f337b50643..5afb425fe421 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpaRagged.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceSdpaRagged.hpp @@ -276,7 +276,8 @@ class CpuFpReferenceSdpaRagged } // Same per-batch window mask as the kernel. Keeps skv in - // [sq + windowOffset - leftBound, sq + windowOffset + rightBound]. A negative bound is open. + // [sq + windowOffset - leftBound + 1, sq + windowOffset + rightBound] (cuDNN: leftBound = L + // keeps L keys, diagonal included). A negative bound is open. static bool isMasked( int64_t sq, int64_t skv, int64_t leftBound, int64_t rightBound, int64_t windowOffset) { @@ -288,7 +289,7 @@ class CpuFpReferenceSdpaRagged return true; } } - if(leftBound >= 0 && skv < sq + windowOffset - leftBound) + if(leftBound >= 0 && skv <= sq + windowOffset - leftBound) { return true; } diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/cpu_graph_executor/detail/PlanUtils.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/cpu_graph_executor/detail/PlanUtils.hpp index 5a3a28b89819..6a9245063da3 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/cpu_graph_executor/detail/PlanUtils.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/cpu_graph_executor/detail/PlanUtils.hpp @@ -40,6 +40,12 @@ DiagonalBandParams extractDiagonalBandParams(const SdpaAttributesType& nodeAttri std::string(planName) + ": left_bound and right_bound must be >= -1 (got left_bound=" + std::to_string(leftBound) + ", right_bound=" + std::to_string(rightBound) + ")"); } + // left_bound follows cuDNN: L keeps the L keys ending at the diagonal; 0 is rejected. + if(leftBound == 0) + { + throw std::invalid_argument(std::string(planName) + + ": left_bound must be -1 (unbounded) or >= 1 (got 0)"); + } const bool isTopLeft = nodeAttributes.diagonal_alignment() == hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment::TOP_LEFT; diff --git a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceSdpa.cpp b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceSdpa.cpp index aaf4a1dc649e..1255d479c3ff 100644 --- a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceSdpa.cpp +++ b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceSdpa.cpp @@ -561,8 +561,9 @@ TEST(TestCpuFpReferenceSdpaFp64, BottomRightCausalMaskLargerSq) TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowTopLeftSquare) { - // Generic sliding window with leftBound=1, rightBound=1 (TopLeft alignment). - // Square shape Sq=Skv=4. Window is centered at sq, width 3. + // Generic sliding window with leftBound=2, rightBound=1 (TopLeft alignment). + // leftBound=2 keeps 2 keys ending at the diagonal (cuDNN), so with rightBound=1 + // the window is centered at sq, width 3. Square shape Sq=Skv=4. // Q=K=0 → all unmasked scores equal → softmax is uniform over the window. // V = {1, 2, 3, 4} → output is the mean of unmasked V values. // @@ -593,7 +594,7 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowTopLeftSquare) o, std::nullopt, noMask, - /*leftBound=*/1, + /*leftBound=*/2, /*rightBound=*/1, /*topLeftAlignment=*/true); @@ -605,8 +606,9 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowTopLeftSquare) TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricTopLeft) { - // Generic sliding window with leftBound=2, rightBound=1 (TopLeft alignment). + // Generic sliding window with leftBound=3, rightBound=1 (TopLeft alignment). // Asymmetric window (wider on the left) with Sq=4, Skv=5. + // Unmasked iff (sq - leftBound) < skv <= (sq + rightBound) // // Expected mask (1=unmasked): // _____kv______ @@ -636,7 +638,7 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricTopLeft) o, std::nullopt, noMask, - /*leftBound=*/2, + /*leftBound=*/3, /*rightBound=*/1, /*topLeftAlignment=*/true); @@ -648,10 +650,10 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricTopLeft) TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricBottomRight) { - // Generic sliding window with leftBound=2, rightBound=1 (BottomRight alignment). + // Generic sliding window with leftBound=3, rightBound=1 (BottomRight alignment). // Asymmetric window (wider on the left) with Sq=4, Skv=5. // diagonal offset = Skv - Sq = 1. - // Unmasked iff (sq + offset - leftBound) <= skv <= (sq + offset + rightBound) + // Unmasked iff (sq + offset - leftBound) < skv <= (sq + offset + rightBound) // iff (sq - 1) <= skv <= (sq + 2) // // Expected mask (1=unmasked): @@ -682,7 +684,7 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricBottomRight) o, std::nullopt, noMask, - /*leftBound=*/2, + /*leftBound=*/3, /*rightBound=*/1, /*topLeftAlignment=*/false); @@ -694,9 +696,9 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricBottomRight) TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSkv) { - // Generic sliding window with leftBound=1, rightBound=1, BottomRight alignment. + // Generic sliding window with leftBound=2, rightBound=1, BottomRight alignment. // Sq=2, Skv=4 → diagonal offset = Skv - Sq = 2. - // Unmasked iff (sq + offset - leftBound) <= skv <= (sq + offset + rightBound) + // Unmasked iff (sq + offset - leftBound) < skv <= (sq + offset + rightBound) // iff (sq + 1) <= skv <= (sq + 3) // // Expected mask (1=unmasked): @@ -724,7 +726,7 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSkv) o, std::nullopt, noMask, - /*leftBound=*/1, + /*leftBound=*/2, /*rightBound=*/1, /*topLeftAlignment=*/false); @@ -734,9 +736,9 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSkv) TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSq) { - // Generic sliding window with leftBound=1, rightBound=1, BottomRight alignment. + // Generic sliding window with leftBound=2, rightBound=1, BottomRight alignment. // Sq=4, Skv=2 → diagonal offset = Skv - Sq = -2. - // Unmasked iff (sq - 2 - 1) <= skv <= (sq - 2 + 1) + // Unmasked iff (sq - 2 - 2) < skv <= (sq - 2 + 1) // iff (sq - 3) <= skv <= (sq - 1) // // Expected mask (1=unmasked): @@ -764,7 +766,7 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSq) o, std::nullopt, noMask, - /*leftBound=*/1, + /*leftBound=*/2, /*rightBound=*/1, /*topLeftAlignment=*/false); @@ -774,6 +776,60 @@ TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSq) EXPECT_NEAR(o.getHostValue(0, 0, 3, 0), 1.5, 1e-5); } +TEST(TestCpuFpReferenceSdpaFp64, LeftBoundKeepsLKeysIncludingDiagonalLikeCudnn) +{ + // cuDNN's left bound: leftBound = L keeps the L keys ending at the diagonal, + // i.e. skv > sq - L (cudnn-frontend sliding_window_mask keeps col + L > row). + // With L = 2, TopLeft, query 5 attends to keys {4, 5}, not {3, 4, 5}. + // Q=K=0 makes the softmax uniform over the kept keys; V = {1..6}. + constexpr int64_t SEQ = 6; + Tensor q({1, 1, SEQ, 1}); + Tensor k({1, 1, SEQ, 1}); + Tensor v({1, 1, SEQ, 1}); + Tensor o({1, 1, SEQ, 1}); + + q.fillWithValue(0.0); + k.fillWithValue(0.0); + for(int64_t s = 0; s < SEQ; ++s) + { + v.setHostValue(static_cast(s + 1), 0, 0, s, 0); + } + + const TensorBase* noMask = nullptr; + CpuFpReferenceSdpa::forward(q, + k, + v, + o, + std::nullopt, + noMask, + /*leftBound=*/2, + /*rightBound=*/0, + /*topLeftAlignment=*/true); + + EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 1.0, 1e-5); // keys {0} + for(int64_t s = 1; s < SEQ; ++s) + { + // keys {s - 1, s} -> mean(s, s + 1) + EXPECT_NEAR(o.getHostValue(0, 0, s, 0), static_cast(s) + 0.5, 1e-5) + << "query " << s; + } + + // L = 1 keeps only the diagonal. + CpuFpReferenceSdpa::forward(q, + k, + v, + o, + std::nullopt, + noMask, + /*leftBound=*/1, + /*rightBound=*/0, + /*topLeftAlignment=*/true); + for(int64_t s = 0; s < SEQ; ++s) + { + EXPECT_NEAR(o.getHostValue(0, 0, s, 0), static_cast(s + 1), 1e-5) << "query " << s; + } +} + TEST(TestCpuFpReferenceSdpaFp64, CausalMaskFutureTokensHaveNoEffect) { // Verify the causal property: changing V values at masked (future) kv positions diff --git a/projects/hipdnn/test_sdk/tests/utilities/cpu_graph_executor/TestSdpaFwdPlan.cpp b/projects/hipdnn/test_sdk/tests/utilities/cpu_graph_executor/TestSdpaFwdPlan.cpp index 671d4414ae93..8654eff8a374 100644 --- a/projects/hipdnn/test_sdk/tests/utilities/cpu_graph_executor/TestSdpaFwdPlan.cpp +++ b/projects/hipdnn/test_sdk/tests/utilities/cpu_graph_executor/TestSdpaFwdPlan.cpp @@ -1227,3 +1227,21 @@ TEST(TestSdpaFwdPlanMask, BuilderAcceptsOnlyMatchingMaskDtype) EXPECT_FALSE( noMaskBuilder.isApplicable(floatMaskWrapper.getNode(0), floatMaskWrapper.getTensorMap())); } + +// left_bound follows cuDNN: L keeps the L keys ending at the diagonal, so 0 keeps no key at or +// before it, and the plans reject it. The frontend rejects it too, so the attributes are built +// directly. +TEST(TestSdpaDiagonalBandParams, RejectsZeroLeftBound) +{ + const auto bandFor = [](std::optional leftBound) { + SdpaAttributesT attrs; + attrs.left_bound = leftBound; + flatbuffers::FlatBufferBuilder builder; + builder.Finish(CreateSdpaAttributes(builder, &attrs)); + const auto* parsed = flatbuffers::GetRoot(builder.GetBufferPointer()); + return extractDiagonalBandParams(*parsed, "SdpaFwdPlan"); + }; + EXPECT_THROW(bandFor(0), std::invalid_argument); + EXPECT_EQ(bandFor(1).leftBound, 1); + EXPECT_EQ(bandFor(std::nullopt).leftBound, -1); +} diff --git a/projects/hipdnn/tests/frontend/IntegrationSdpaBwdDescriptorLowering.cpp b/projects/hipdnn/tests/frontend/IntegrationSdpaBwdDescriptorLowering.cpp index 5ddbb46a2245..32529473e029 100644 --- a/projects/hipdnn/tests/frontend/IntegrationSdpaBwdDescriptorLowering.cpp +++ b/projects/hipdnn/tests/frontend/IntegrationSdpaBwdDescriptorLowering.cpp @@ -399,7 +399,7 @@ TEST_F(IntegrationSdpaBwdDescriptorLowering, SdpaBwdWithAllOptionalTensorsAndSca sdpaAttrs.causal_mask = true; sdpaAttrs.causal_mask_bottom_right = true; sdpaAttrs.set_attn_scale(0.125f); - sdpaAttrs.set_diagonal_band_left_bound(0); + sdpaAttrs.set_diagonal_band_left_bound(64); sdpaAttrs.set_diagonal_band_right_bound(128); sdpaAttrs.set_diagonal_alignment(DiagonalAlignment::BOTTOM_RIGHT); @@ -587,7 +587,7 @@ TEST_F(IntegrationSdpaBwdDescriptorLowering, SdpaBwdWithAllOptionalTensorsAndSca // Integer scalars ASSERT_TRUE(sdpa->left_bound.has_value()); - EXPECT_EQ(sdpa->left_bound.value(), 0); + EXPECT_EQ(sdpa->left_bound.value(), 64); ASSERT_TRUE(sdpa->right_bound.has_value()); EXPECT_EQ(sdpa->right_bound.value(), 128); diff --git a/projects/hipdnn/tests/frontend/IntegrationSdpaFwdDescriptorLowering.cpp b/projects/hipdnn/tests/frontend/IntegrationSdpaFwdDescriptorLowering.cpp index ac2aed8e4d09..614ca2b8265a 100644 --- a/projects/hipdnn/tests/frontend/IntegrationSdpaFwdDescriptorLowering.cpp +++ b/projects/hipdnn/tests/frontend/IntegrationSdpaFwdDescriptorLowering.cpp @@ -521,7 +521,7 @@ TEST_F(IntegrationSdpaFwdDescriptorLowering, SdpaFwdWithOptionalTensorsAndScalar sdpaAttrs.set_generate_stats(true); sdpaAttrs.causal_mask = true; sdpaAttrs.set_attn_scale(0.125f); - sdpaAttrs.set_diagonal_band_left_bound(0); + sdpaAttrs.set_diagonal_band_left_bound(64); sdpaAttrs.set_diagonal_band_right_bound(128); sdpaAttrs.set_diagonal_alignment(DiagonalAlignment::BOTTOM_RIGHT); sdpaAttrs.set_implementation(AttentionImplementation::UNIFIED); @@ -591,7 +591,7 @@ TEST_F(IntegrationSdpaFwdDescriptorLowering, SdpaFwdWithOptionalTensorsAndScalar // Integer scalars ASSERT_TRUE(sdpa->left_bound.has_value()); - EXPECT_EQ(sdpa->left_bound.value(), 0); + EXPECT_EQ(sdpa->left_bound.value(), 64); ASSERT_TRUE(sdpa->right_bound.has_value()); EXPECT_EQ(sdpa->right_bound.value(), 128); diff --git a/projects/hipdnn/tools/IngestorGenerator/tests/test_mine_shapes.py b/projects/hipdnn/tools/IngestorGenerator/tests/test_mine_shapes.py index d3a8c65cdd9f..085553e96caa 100644 --- a/projects/hipdnn/tools/IngestorGenerator/tests/test_mine_shapes.py +++ b/projects/hipdnn/tools/IngestorGenerator/tests/test_mine_shapes.py @@ -393,7 +393,7 @@ def test_the_deprecated_boolean_with_bottom_right_and_square_shape_is_causal( def test_the_deprecated_boolean_keeps_a_stored_left_bound(self, tmp_path): self._graph( tmp_path / "graphs" / "g.json", - {"causal_mask": True, "left_bound": 127, "right_bound": None}, + {"causal_mask": True, "left_bound": 128, "right_bound": None}, ) shapes = self._mine(tmp_path) assert [ @@ -429,6 +429,14 @@ def test_a_non_numeric_left_bound_is_refused_not_resolved_to_causal(self, tmp_pa ) self._assert_refused(tmp_path) + def test_a_zero_left_bound_is_refused(self, tmp_path): + """left_bound follows cuDNN (L keys ending at the diagonal); 0 is refused.""" + self._graph( + tmp_path / "graphs" / "g.json", + {"causal_mask": False, "left_bound": 0, "right_bound": 0}, + ) + self._assert_refused(tmp_path) + class TestRocKeBenchTree: """The third source: rocKE's own benchmark tree, the only one saying what the kernel @@ -687,7 +695,7 @@ def _mine(tmp_path: Path) -> tuple[int, str, list]: _UNMASKED = {"causal_mask": False, "left_bound": None, "right_bound": None} _CAUSAL = {"causal_mask": False, "left_bound": -1, "right_bound": 0} - _WINDOW = {"causal_mask": False, "left_bound": 127, "right_bound": 0} + _WINDOW = {"causal_mask": False, "left_bound": 128, "right_bound": 0} def test_an_asymmetric_v_head_dimension_is_carried_not_copied_from_q( self, tmp_path @@ -748,11 +756,11 @@ def test_two_windows_of_different_width_do_not_merge(self, tmp_path): as the same request.""" self._graph( tmp_path / "graphs" / "w64.json", - attrs={"causal_mask": False, "left_bound": 63, "right_bound": 0}, + attrs={"causal_mask": False, "left_bound": 64, "right_bound": 0}, ) self._graph( tmp_path / "graphs" / "w128.json", - attrs={"causal_mask": False, "left_bound": 127, "right_bound": 0}, + attrs={"causal_mask": False, "left_bound": 128, "right_bound": 0}, ) rc, output, shapes = self._mine(tmp_path) assert rc == 0, output @@ -1456,7 +1464,7 @@ def test_the_window_kind_has_no_anchor_to_carry_it(self, tmp_path): tmp_path / "graphs" / "g.json", 512, 4096, - {"left_bound": 127, "right_bound": 0, "diagonal_alignment": "BOTTOM_RIGHT"}, + {"left_bound": 128, "right_bound": 0, "diagonal_alignment": "BOTTOM_RIGHT"}, ) out = tmp_path / "shapes.json" result = subprocess.run( @@ -1531,7 +1539,7 @@ def test_every_mined_mask_resolves_to_the_request_it_describes( {"left_bound": -1, "right_bound": 0, "diagonal_alignment": "TOP_LEFT"}, ) _sdpa_graph( - graphs / "swa.json", 4096, 4096, {"left_bound": 127, "right_bound": 0} + graphs / "swa.json", 4096, 4096, {"left_bound": 128, "right_bound": 0} ) catalog = tmp_path / "MODEL_CATALOG.md" catalog.write_text(_CATALOG) @@ -1557,7 +1565,7 @@ def test_every_mined_mask_resolves_to_the_request_it_describes( "prefill": {frozenset({"causal"})}, "decode": {frozenset({"causal", "causal_bottom_right"})}, } - # left_bound 127, right_bound 0: the request states the width and a right bound + # left_bound 128, right_bound 0: the request states the width and a right bound # of 0, which rocKE normalizes to a band 128 keys wide ending on the diagonal. assert [ ( diff --git a/projects/hipdnn/tools/IngestorGenerator/tools/mine_shapes.py b/projects/hipdnn/tools/IngestorGenerator/tools/mine_shapes.py index 8d4c78b8c176..55bee6797cb4 100644 --- a/projects/hipdnn/tools/IngestorGenerator/tools/mine_shapes.py +++ b/projects/hipdnn/tools/IngestorGenerator/tools/mine_shapes.py @@ -175,6 +175,9 @@ def _mask_from_attributes( raise SystemExit( f"FAIL: {path}: invalid {name} {value!r}; expected null or integer >= -1" ) + # left_bound follows cuDNN (L keeps the L keys ending at the diagonal); 0 is refused. + if left == 0: + raise SystemExit(f"FAIL: {path}: invalid left_bound 0; expected -1 or >= 1") causal = attrs.get("causal_mask") causal_bottom_right = attrs.get("causal_mask_bottom_right") if causal and causal_bottom_right: @@ -202,7 +205,9 @@ def _mask_from_attributes( "mask_type": _causal_mask_type(bottom_right, seqlen_q, seqlen_k), "sliding_window": 0, } - return _mask_fields(SLIDING_WINDOW, left + 1) + # hipDNN's left_bound and the request's sliding_window both count the keys kept, + # diagonal included (cuDNN semantics), so the width passes through unchanged. + return _mask_fields(SLIDING_WINDOW, left) #: Every source spelling of a dtype -> the rocKE spec's spelling. An unmapped dtype