derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
2488 lines
92 KiB
C++
2488 lines
92 KiB
C++
// Copyright © Advanced Micro Devices, Inc., or its affiliates.
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
#include <cmath>
|
||
#include <functional>
|
||
#include <gtest/gtest.h>
|
||
#include <hipdnn_data_sdk/types.hpp>
|
||
#include <hipdnn_data_sdk/utilities/Tensor.hpp>
|
||
#include <hipdnn_test_sdk/utilities/CpuFpReferenceSdpa.hpp>
|
||
#include <hipdnn_test_sdk/utilities/detail/CpuFpReferenceUtilities.hpp>
|
||
#include <hipdnn_test_sdk/utilities/detail/GradientCheckUtilities.hpp>
|
||
|
||
using namespace hipdnn_test_sdk::utilities;
|
||
using namespace hipdnn_data_sdk::utilities;
|
||
using namespace hipdnn_data_sdk::types;
|
||
using hipdnn_test_sdk::detail::safeTestTypeCast;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TypePair helper for TYPED_TEST
|
||
// ---------------------------------------------------------------------------
|
||
|
||
template <typename T1, typename T2>
|
||
struct TypePair
|
||
{
|
||
using First = T1;
|
||
using Second = T2;
|
||
};
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Fp64 precision unit tests
|
||
// ---------------------------------------------------------------------------
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, SanityCheck)
|
||
{
|
||
// [B=1, H=1, Sq=2, Skv=2, D=2, Dv=2]
|
||
// Q[0,0,0,:] = [1,0], Q[0,0,1,:] = [0,1]
|
||
// K[0,0,0,:] = [1,0], K[0,0,1,:] = [0,1]
|
||
// V[0,0,0,:] = [1,2], V[0,0,1,:] = [3,4]
|
||
// Default scale = 1/sqrt(2)
|
||
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 1, 2, 2});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
q.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 0, 1);
|
||
v.setHostValue(3.0, 0, 0, 1, 0);
|
||
v.setHostValue(4.0, 0, 0, 1, 1);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Expected using default scale = 1/sqrt(2):
|
||
// sq=0: S[0]=scale, S[1]=0 → P[0]=1/(1+exp(-scale)), P[1]=exp(-scale)/(1+exp(-scale))
|
||
// sq=1: S[0]=0, S[1]=scale → P[0]=exp(-scale)/(1+exp(-scale)), P[1]=1/(1+exp(-scale))
|
||
const float scale = 1.0f / std::sqrt(2.0f);
|
||
const float eLow = std::exp(-scale);
|
||
const float sumExp = 1.0f + eLow;
|
||
const float pHigh = 1.0f / sumExp; // weight for the matching kv token
|
||
const float pLow = eLow / sumExp; // weight for the non-matching kv token
|
||
|
||
const double tol = 1e-5;
|
||
|
||
// sq=0 biased toward kv=0 (V=[1,2])
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), static_cast<double>(pHigh * 1.0f + pLow * 3.0f), tol);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 1), static_cast<double>(pHigh * 2.0f + pLow * 4.0f), tol);
|
||
|
||
// sq=1 biased toward kv=1 (V=[3,4])
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), static_cast<double>(pLow * 1.0f + pHigh * 3.0f), tol);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 1), static_cast<double>(pLow * 2.0f + pHigh * 4.0f), tol);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, DefaultScaleIs1OverSqrtD)
|
||
{
|
||
// Verify that the default attention scale equals 1/sqrt(headDim).
|
||
// Q[0,0,0,0]=1, rest zero; K[0,0,0,0]=1 (dot=1), K[0,0,1,:]=0 (dot=0).
|
||
const int64_t headDim = 4;
|
||
|
||
Tensor<double> q({1, 1, 1, headDim});
|
||
Tensor<double> k({1, 1, 2, headDim});
|
||
Tensor<double> v({1, 1, 2, 1});
|
||
Tensor<double> o({1, 1, 1, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
v.fillWithValue(0.0);
|
||
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Default scale = 1/sqrt(headDim)
|
||
// S[0] = 1 * scale, S[1] = 0
|
||
// P[0] = 1/(1+exp(-scale)), P[1] = exp(-scale)/(1+exp(-scale))
|
||
const float defaultScale = 1.0f / std::sqrt(static_cast<float>(headDim));
|
||
const float e1 = std::exp(-defaultScale);
|
||
const float sumE = 1.0f + e1;
|
||
const float p0 = 1.0f / sumE;
|
||
const float p1 = e1 / sumE;
|
||
const auto expected = static_cast<double>(p0 * 1.0f + p1 * 2.0f);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), expected, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, CustomScale)
|
||
{
|
||
// Verify that an explicit attnScaleValue overrides the default 1/sqrt(D).
|
||
const int64_t headDim = 4;
|
||
|
||
Tensor<double> q({1, 1, 1, headDim});
|
||
Tensor<double> k({1, 1, 2, headDim});
|
||
Tensor<double> v({1, 1, 2, 1});
|
||
Tensor<double> oDefault({1, 1, 1, 1});
|
||
Tensor<double> oCustom({1, 1, 1, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
v.fillWithValue(0.0);
|
||
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, oDefault); // default scale = 0.5
|
||
CpuFpReferenceSdpa::forward(q, k, v, oCustom, std::optional<float>{2.0f}); // custom scale = 2.0
|
||
|
||
// Expected with custom scale = 2.0: S[0]=2, S[1]=0
|
||
const float customScale = 2.0f;
|
||
const float e1Custom = std::exp(-customScale);
|
||
const float sumCustom = 1.0f + e1Custom;
|
||
const float p0Custom = 1.0f / sumCustom;
|
||
const float p1Custom = e1Custom / sumCustom;
|
||
const auto expectedCustom = static_cast<double>(p0Custom * 1.0f + p1Custom * 2.0f);
|
||
|
||
EXPECT_NEAR(oCustom.getHostValue(0, 0, 0, 0), expectedCustom, 1e-5);
|
||
|
||
// Custom scale result must differ from default scale result
|
||
EXPECT_NE(oDefault.getHostValue(0, 0, 0, 0), oCustom.getHostValue(0, 0, 0, 0));
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, WithAttnMask)
|
||
{
|
||
// Same Q/K/V as SanityCheck, but with an additive mask that suppresses kv=1 for sq=0.
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 1, 2, 2});
|
||
Tensor<float> mask({2, 2}); // [Sq, Skv]
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
q.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 0, 1);
|
||
v.setHostValue(3.0, 0, 0, 1, 0);
|
||
v.setHostValue(4.0, 0, 0, 1, 1);
|
||
|
||
// For sq=0: block kv=1 with a large negative value
|
||
mask.setHostValue(0.0f, 0, 0);
|
||
mask.setHostValue(-1e4f, 0, 1);
|
||
// For sq=1: no masking
|
||
mask.setHostValue(0.0f, 1, 0);
|
||
mask.setHostValue(0.0f, 1, 1);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, &mask);
|
||
|
||
// sq=0: kv=1 is masked → O ≈ V[0,0,0,:] = [1, 2]
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 1.0, 1e-3);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 1), 2.0, 1e-3);
|
||
|
||
// sq=1: unmasked, biased toward kv=1 (higher score) → O closer to V[0,0,1,:]=[3,4]
|
||
EXPECT_GT(o.getHostValue(0, 0, 1, 0), 2.0);
|
||
EXPECT_GT(o.getHostValue(0, 0, 1, 1), 3.0);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, GqaSupport)
|
||
{
|
||
// H=4, Hkv=2: each KV head serves 2 Q heads.
|
||
// Skv=1: softmax is trivially 1.0, so O = V[b, kvHead, 0, :].
|
||
Tensor<double> q({1, 4, 1, 2});
|
||
Tensor<double> k({1, 2, 1, 2});
|
||
Tensor<double> v({1, 2, 1, 2});
|
||
Tensor<double> o({1, 4, 1, 2});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 0, 1); // V[kvHead=0] = [1, 2]
|
||
v.setHostValue(3.0, 0, 1, 0, 0);
|
||
v.setHostValue(4.0, 0, 1, 0, 1); // V[kvHead=1] = [3, 4]
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
const double tol = 1e-5;
|
||
|
||
// h=0,1: kvHead = h/2 = 0 → O = V[0,0,0,:] = [1, 2]
|
||
for(int64_t h = 0; h < 2; ++h)
|
||
{
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 0), 1.0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 1), 2.0, tol);
|
||
}
|
||
|
||
// h=2,3: kvHead = h/2 = 1 → O = V[0,1,0,:] = [3, 4]
|
||
for(int64_t h = 2; h < 4; ++h)
|
||
{
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 0), 3.0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 1), 4.0, tol);
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, MqaSupport)
|
||
{
|
||
// H=4, Hkv=1: single KV head for all Q heads.
|
||
// Skv=1: O = V[0, 0, 0, :] for all Q heads.
|
||
Tensor<double> q({1, 4, 1, 2});
|
||
Tensor<double> k({1, 1, 1, 2});
|
||
Tensor<double> v({1, 1, 1, 2});
|
||
Tensor<double> o({1, 4, 1, 2});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(5.0, 0, 0, 0, 0);
|
||
v.setHostValue(6.0, 0, 0, 0, 1); // V[kvHead=0] = [5, 6]
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
const double tol = 1e-5;
|
||
|
||
for(int64_t h = 0; h < 4; ++h)
|
||
{
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 0), 5.0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 1), 6.0, tol);
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, GqaDifferentKVHeads)
|
||
{
|
||
// H=4, Hk=2, Hv=1, Skv=2: K has 2 heads, V has 1 head.
|
||
// With Skv=2, distinct K-head values produce different attention patterns,
|
||
// verifying both K-head and V-head mappings.
|
||
// K head 0: [1,0] at skv=0, [0,0] at skv=1 → Q·K biased toward skv=0
|
||
// K head 1: [0,0] at skv=0, [1,0] at skv=1 → Q·K biased toward skv=1
|
||
// V head 0: [1,2] at skv=0, [3,4] at skv=1
|
||
// Q heads 0,1 → K head 0 (biased toward V[skv=0]=[1,2])
|
||
// Q heads 2,3 → K head 1 (biased toward V[skv=1]=[3,4])
|
||
Tensor<double> q({1, 4, 1, 2});
|
||
Tensor<double> k({1, 2, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 4, 1, 2});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
// Q: all heads have query [1, 0]
|
||
for(int64_t h = 0; h < 4; ++h)
|
||
{
|
||
q.setHostValue(1.0, 0, h, 0, 0);
|
||
}
|
||
|
||
// K head 0: [1,0] at skv=0, [0,0] at skv=1
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
// K head 1: [0,0] at skv=0, [1,0] at skv=1
|
||
k.setHostValue(1.0, 0, 1, 1, 0);
|
||
|
||
// V head 0: [1,2] at skv=0, [3,4] at skv=1
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 0, 1);
|
||
v.setHostValue(3.0, 0, 0, 1, 0);
|
||
v.setHostValue(4.0, 0, 0, 1, 1);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
const float scale = 1.0f / std::sqrt(2.0f);
|
||
const float eLow = std::exp(-scale);
|
||
const float sumExp = 1.0f + eLow;
|
||
const float pHigh = 1.0f / sumExp;
|
||
const float pLow = eLow / sumExp;
|
||
const double tol = 1e-5;
|
||
|
||
// Q heads 0,1 → K head 0 → biased toward skv=0: O ≈ pHigh*V[0] + pLow*V[1]
|
||
const auto expectedK0D0 = static_cast<double>(pHigh * 1.0f + pLow * 3.0f);
|
||
const auto expectedK0D1 = static_cast<double>(pHigh * 2.0f + pLow * 4.0f);
|
||
for(int64_t h = 0; h < 2; ++h)
|
||
{
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 0), expectedK0D0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 1), expectedK0D1, tol);
|
||
}
|
||
|
||
// Q heads 2,3 → K head 1 → biased toward skv=1: O ≈ pLow*V[0] + pHigh*V[1]
|
||
const auto expectedK1D0 = static_cast<double>(pLow * 1.0f + pHigh * 3.0f);
|
||
const auto expectedK1D1 = static_cast<double>(pLow * 2.0f + pHigh * 4.0f);
|
||
for(int64_t h = 2; h < 4; ++h)
|
||
{
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 0), expectedK1D0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, h, 0, 1), expectedK1D1, tol);
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, AttnMaskBroadcastRank2)
|
||
{
|
||
// Rank-2 mask [1, Skv]: dim[0]=1 broadcasts over all sq (and also over batch and head
|
||
// which are absent from the rank-2 mask).
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 1, 2, 2});
|
||
Tensor<float> mask({1, 2}); // [1, Skv]
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
q.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 0, 1);
|
||
v.setHostValue(3.0, 0, 0, 1, 0);
|
||
v.setHostValue(4.0, 0, 0, 1, 1);
|
||
|
||
// Block kv=1 for ALL sq positions (mask broadcasts over sq via dim[0]=1)
|
||
mask.setHostValue(0.0f, 0, 0);
|
||
mask.setHostValue(-1e4f, 0, 1);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, &mask);
|
||
|
||
// Both sq=0 and sq=1 should have O ≈ V[0,0,0,:] = [1, 2]
|
||
const double tol = 1e-3;
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 1.0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 1), 2.0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), 1.0, tol);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 1), 2.0, tol);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, CausalMask)
|
||
{
|
||
// causalMask=true: sq can attend to kv <= sq only.
|
||
// [B=1, H=1, Sq=3, Skv=3, D=3, Dv=1]
|
||
// Q[sq,:] = one-hot(sq), K[skv,:] = one-hot(skv) → dot(Q[sq], K[skv]) = δ(sq,skv)
|
||
// V[skv,0] = skv+1: [1, 2, 3]
|
||
// scale = 1/sqrt(3), eNeg = exp(-scale)
|
||
//
|
||
// sq=0: only kv=0 unmasked, S[0]=scale → P[0]≈1 → O≈1
|
||
// sq=1: kv=0,1 unmasked, S[0]=0, S[1]=scale → O = eNeg/(1+eNeg)*1 + 1/(1+eNeg)*2
|
||
// sq=2: all unmasked, S[0]=S[1]=0, S[2]=scale → O = (eNeg+2*eNeg+3)/(2*eNeg+1)
|
||
|
||
Tensor<double> q({1, 1, 3, 3});
|
||
Tensor<double> k({1, 1, 3, 3});
|
||
Tensor<double> v({1, 1, 3, 1});
|
||
Tensor<double> o({1, 1, 3, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
// Q[sq,:] = one-hot(sq)
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
q.setHostValue(1.0, 0, 0, 1, 1);
|
||
q.setHostValue(1.0, 0, 0, 2, 2);
|
||
|
||
// K[skv,:] = one-hot(skv)
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 1, 1);
|
||
k.setHostValue(1.0, 0, 0, 2, 2);
|
||
|
||
// V[skv,0] = skv+1
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
v.setHostValue(3.0, 0, 0, 2, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, /*causalMask=*/true);
|
||
|
||
const float scale = 1.0f / std::sqrt(3.0f);
|
||
const float eNeg = std::exp(-scale);
|
||
|
||
// sq=0: only kv=0 unmasked → P[0]≈1 → O≈V[0,0]=1
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 1.0, 1e-3);
|
||
|
||
// sq=1: kv=0,1 unmasked
|
||
// S[0]=0, S[1]=scale → maxVal=scale
|
||
// P[0]=eNeg/(1+eNeg), P[1]=1/(1+eNeg)
|
||
const float p0sq1 = eNeg / (1.0f + eNeg);
|
||
const float p1sq1 = 1.0f / (1.0f + eNeg);
|
||
const auto expSq1 = static_cast<double>(p0sq1 * 1.0f + p1sq1 * 2.0f);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), expSq1, 1e-5);
|
||
|
||
// sq=2: all kv unmasked
|
||
// S[0]=S[1]=0, S[2]=scale → maxVal=scale
|
||
// P[0]=P[1]=eNeg/(2*eNeg+1), P[2]=1/(2*eNeg+1)
|
||
const float denom = 2.0f * eNeg + 1.0f;
|
||
const auto expSq2 = static_cast<double>((eNeg * 1.0f + eNeg * 2.0f + 1.0f * 3.0f) / denom);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 2, 0), expSq2, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, BottomRightCausalMaskLargerSkv)
|
||
{
|
||
// Causal mask (leftBound=-1, rightBound=0) with BottomRight alignment.
|
||
// Sq=2, Skv=4 → diagonal offset = Skv - Sq = 2.
|
||
// Unmasked iff skv <= sq + offset + rightBound = sq + 2.
|
||
// Q=K=0 → all unmasked scores equal → softmax is uniform over the unmasked set.
|
||
// V = {1, 2, 3, 4} → output is the mean of unmasked V values.
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// ___kv____
|
||
// q | 1 1 1 0 sq=0: skv ∈ [-∞, 2] → {0, 1, 2} → mean(1,2,3) = 6/3 = 2.0
|
||
// | 1 1 1 1 sq=1: skv ∈ [-∞, 3] → {0, 1, 2, 3} → mean(1,2,3,4) = 10/4 = 2.5
|
||
|
||
Tensor<double> q({1, 1, 2, 1});
|
||
Tensor<double> k({1, 1, 4, 1});
|
||
Tensor<double> v({1, 1, 4, 1});
|
||
Tensor<double> o({1, 1, 2, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
v.setHostValue(3.0, 0, 0, 2, 0);
|
||
v.setHostValue(4.0, 0, 0, 3, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/-1,
|
||
/*rightBound=*/0,
|
||
/*topLeftAlignment=*/false);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 6.0 / 3.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), 10.0 / 4.0, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, BottomRightCausalMaskLargerSq)
|
||
{
|
||
// Causal mask (leftBound=-1, rightBound=0) with BottomRight alignment.
|
||
// Sq=4, Skv=2 → diagonal offset = Skv - Sq = -2.
|
||
// Unmasked iff skv <= sq + offset + rightBound = sq - 2.
|
||
// Q=K=0 → all unmasked scores equal → softmax is uniform over the unmasked set.
|
||
// V = {1, 2} → output is the mean of unmasked V values.
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// _kv__
|
||
// | 0 0 sq=0: skv ∈ [-∞, -2] → {} → fully masked, output = 0
|
||
// q | 0 0 sq=1: skv ∈ [-∞, -1] → {} → fully masked, output = 0
|
||
// | 1 0 sq=2: skv ∈ [-∞, 0] → {0} → mean(1) = 1.0
|
||
// | 1 1 sq=3: skv ∈ [-∞, 1] → {0, 1} → mean(1,2) = 1.5
|
||
|
||
Tensor<double> q({1, 1, 4, 1});
|
||
Tensor<double> k({1, 1, 2, 1});
|
||
Tensor<double> v({1, 1, 2, 1});
|
||
Tensor<double> o({1, 1, 4, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/-1,
|
||
/*rightBound=*/0,
|
||
/*topLeftAlignment=*/false);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 0.0, 1e-5); // fully masked
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), 0.0, 1e-5); // fully masked
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 2, 0), 1.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 3, 0), 1.5, 1e-5);
|
||
}
|
||
|
||
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.
|
||
// 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.
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// ___kv_____
|
||
// q | 1 1 0 0 sq=0: skv ∈ [-1, 1] → {0, 1} → mean(1,2) = 1.5
|
||
// | 1 1 1 0 sq=1: skv ∈ [ 0, 2] → {0, 1, 2} → mean(1,2,3) = 2.0
|
||
// | 0 1 1 1 sq=2: skv ∈ [ 1, 3] → {1, 2, 3} → mean(2,3,4) = 3.0
|
||
// | 0 0 1 1 sq=3: skv ∈ [ 2, 4] → {2, 3} → mean(3,4) = 3.5
|
||
|
||
Tensor<double> q({1, 1, 4, 1});
|
||
Tensor<double> k({1, 1, 4, 1});
|
||
Tensor<double> v({1, 1, 4, 1});
|
||
Tensor<double> o({1, 1, 4, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
v.setHostValue(3.0, 0, 0, 2, 0);
|
||
v.setHostValue(4.0, 0, 0, 3, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/1,
|
||
/*rightBound=*/1,
|
||
/*topLeftAlignment=*/true);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), (1.0 + 2.0) / 2.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), (1.0 + 2.0 + 3.0) / 3.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 2, 0), (2.0 + 3.0 + 4.0) / 3.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 3, 0), (3.0 + 4.0) / 2.0, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricTopLeft)
|
||
{
|
||
// Generic sliding window with leftBound=2, rightBound=1 (TopLeft alignment).
|
||
// Asymmetric window (wider on the left) with Sq=4, Skv=5.
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// _____kv______
|
||
// q | 1 1 0 0 0 sq=0: skv ∈ [-2, 1] → {0, 1} → mean(1,2) = 1.5
|
||
// | 1 1 1 0 0 sq=1: skv ∈ [-1, 2] → {0, 1, 2} → mean(1,2,3) = 2.0
|
||
// | 1 1 1 1 0 sq=2: skv ∈ [ 0, 3] → {0, 1, 2, 3} → mean(1,2,3,4) = 2.5
|
||
// | 0 1 1 1 1 sq=3: skv ∈ [ 1, 4] → {1, 2, 3, 4} → mean(2,3,4,5) = 3.5
|
||
|
||
Tensor<double> q({1, 1, 4, 1});
|
||
Tensor<double> k({1, 1, 5, 1});
|
||
Tensor<double> v({1, 1, 5, 1});
|
||
Tensor<double> o({1, 1, 4, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
v.setHostValue(3.0, 0, 0, 2, 0);
|
||
v.setHostValue(4.0, 0, 0, 3, 0);
|
||
v.setHostValue(5.0, 0, 0, 4, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/2,
|
||
/*rightBound=*/1,
|
||
/*topLeftAlignment=*/true);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), (1.0 + 2.0) / 2.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), (1.0 + 2.0 + 3.0) / 3.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 2, 0), (1.0 + 2.0 + 3.0 + 4.0) / 4.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 3, 0), (2.0 + 3.0 + 4.0 + 5.0) / 4.0, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowAsymmetricBottomRight)
|
||
{
|
||
// Generic sliding window with leftBound=2, 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)
|
||
// iff (sq - 1) <= skv <= (sq + 2)
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// _____kv______
|
||
// q | 1 1 1 0 0 sq=0: skv ∈ [-1, 2] → {0, 1, 2} → mean(1,2,3) = 2.0
|
||
// | 1 1 1 1 0 sq=1: skv ∈ [ 0, 3] → {0, 1, 2, 3} → mean(1,2,3,4) = 2.5
|
||
// | 0 1 1 1 1 sq=2: skv ∈ [ 1, 4] → {1, 2, 3, 4} → mean(2,3,4,5) = 3.5
|
||
// | 0 0 1 1 1 sq=3: skv ∈ [ 2, 5] → {2, 3, 4} → mean(3,4,5) = 4.0
|
||
|
||
Tensor<double> q({1, 1, 4, 1});
|
||
Tensor<double> k({1, 1, 5, 1});
|
||
Tensor<double> v({1, 1, 5, 1});
|
||
Tensor<double> o({1, 1, 4, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
v.setHostValue(3.0, 0, 0, 2, 0);
|
||
v.setHostValue(4.0, 0, 0, 3, 0);
|
||
v.setHostValue(5.0, 0, 0, 4, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/2,
|
||
/*rightBound=*/1,
|
||
/*topLeftAlignment=*/false);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), (1.0 + 2.0 + 3.0) / 3.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), (1.0 + 2.0 + 3.0 + 4.0) / 4.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 2, 0), (2.0 + 3.0 + 4.0 + 5.0) / 4.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 3, 0), (3.0 + 4.0 + 5.0) / 3.0, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSkv)
|
||
{
|
||
// Generic sliding window with leftBound=1, rightBound=1, BottomRight alignment.
|
||
// Sq=2, Skv=4 → diagonal offset = Skv - Sq = 2.
|
||
// Unmasked iff (sq + offset - leftBound) <= skv <= (sq + offset + rightBound)
|
||
// iff (sq + 1) <= skv <= (sq + 3)
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// ___kv____
|
||
// q | 0 1 1 1 sq=0: skv ∈ [1, 3] → {1, 2, 3} → mean(2,3,4) = 3.0
|
||
// | 0 0 1 1 sq=1: skv ∈ [2, 4] → {2, 3} → mean(3,4) = 3.5
|
||
|
||
Tensor<double> q({1, 1, 2, 1});
|
||
Tensor<double> k({1, 1, 4, 1});
|
||
Tensor<double> v({1, 1, 4, 1});
|
||
Tensor<double> o({1, 1, 2, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
v.setHostValue(3.0, 0, 0, 2, 0);
|
||
v.setHostValue(4.0, 0, 0, 3, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/1,
|
||
/*rightBound=*/1,
|
||
/*topLeftAlignment=*/false);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), (2.0 + 3.0 + 4.0) / 3.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), (3.0 + 4.0) / 2.0, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, SlidingWindowBottomRightLargerSq)
|
||
{
|
||
// Generic sliding window with leftBound=1, rightBound=1, BottomRight alignment.
|
||
// Sq=4, Skv=2 → diagonal offset = Skv - Sq = -2.
|
||
// Unmasked iff (sq - 2 - 1) <= skv <= (sq - 2 + 1)
|
||
// iff (sq - 3) <= skv <= (sq - 1)
|
||
//
|
||
// Expected mask (1=unmasked):
|
||
// _kv__
|
||
// | 0 0 sq=0: skv ∈ [-3, -1] → {} → fully masked, output = 0
|
||
// q | 1 0 sq=1: skv ∈ [-2, 0] → {0} → mean(1) = 1.0
|
||
// | 1 1 sq=2: skv ∈ [-1, 1] → {0, 1} → mean(1,2) = 1.5
|
||
// | 1 1 sq=3: skv ∈ [ 0, 2] → {0, 1} → mean(1,2) = 1.5
|
||
|
||
Tensor<double> q({1, 1, 4, 1});
|
||
Tensor<double> k({1, 1, 2, 1});
|
||
Tensor<double> v({1, 1, 2, 1});
|
||
Tensor<double> o({1, 1, 4, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 1, 0);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
std::nullopt,
|
||
noMask,
|
||
/*leftBound=*/1,
|
||
/*rightBound=*/1,
|
||
/*topLeftAlignment=*/false);
|
||
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 0, 0), 0.0, 1e-5); // fully masked
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 1, 0), 1.0, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 2, 0), 1.5, 1e-5);
|
||
EXPECT_NEAR(o.getHostValue(0, 0, 3, 0), 1.5, 1e-5);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, CausalMaskFutureTokensHaveNoEffect)
|
||
{
|
||
// Verify the causal property: changing V values at masked (future) kv positions
|
||
// produces identical output, since those positions contribute zero probability.
|
||
// [B=1, H=1, Sq=2, Skv=3, D=2, Dv=2]
|
||
// causalMask=true: sq=0 sees kv=0 only; sq=1 sees kv=0,1; kv=2 is always masked.
|
||
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 3, 2});
|
||
Tensor<double> vBase({1, 1, 3, 2});
|
||
Tensor<double> vAlt({1, 1, 3, 2}); // identical except kv=2 is wildly different
|
||
Tensor<double> oBase({1, 1, 2, 2});
|
||
Tensor<double> oAlt({1, 1, 2, 2});
|
||
|
||
// Uniform Q and K so all unmasked scores are equal (softmax is uniform)
|
||
q.fillWithValue(1.0);
|
||
k.fillWithValue(1.0);
|
||
vBase.fillWithValue(1.0);
|
||
vAlt.fillWithValue(1.0);
|
||
|
||
// Make the always-masked kv=2 position very different in vAlt
|
||
vAlt.setHostValue(999.0, 0, 0, 2, 0);
|
||
vAlt.setHostValue(999.0, 0, 0, 2, 1);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, vBase, oBase, std::nullopt, noMask, /*causalMask=*/true);
|
||
CpuFpReferenceSdpa::forward(q, k, vAlt, oAlt, std::nullopt, noMask, /*causalMask=*/true);
|
||
|
||
// Outputs must be identical since kv=2 is always masked out
|
||
const double tol = 1e-6;
|
||
for(int64_t sq = 0; sq < 2; ++sq)
|
||
{
|
||
for(int64_t dv = 0; dv < 2; ++dv)
|
||
{
|
||
EXPECT_NEAR(oBase.getHostValue(std::vector<int64_t>{0, 0, sq, dv}),
|
||
oAlt.getHostValue(std::vector<int64_t>{0, 0, sq, dv}),
|
||
tol);
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Multi-type smoke test
|
||
// ---------------------------------------------------------------------------
|
||
|
||
using TypesSdpaFwd
|
||
= ::testing::Types<TypePair<float, float>, TypePair<half, float>, TypePair<bfloat16, float>>;
|
||
|
||
template <class T>
|
||
class CpuFpReferenceSdpaFwd : public ::testing::Test
|
||
{
|
||
};
|
||
|
||
TYPED_TEST_SUITE(CpuFpReferenceSdpaFwd, TypesSdpaFwd, );
|
||
|
||
TYPED_TEST(CpuFpReferenceSdpaFwd, BasicFwd)
|
||
{
|
||
using InT = typename TypeParam::First;
|
||
|
||
// Uniform Q/K/V: all dot products are equal → uniform softmax → O = V_val
|
||
Tensor<InT> q({1, 2, 4, 8});
|
||
Tensor<InT> k({1, 2, 4, 8});
|
||
Tensor<InT> v({1, 2, 4, 8});
|
||
Tensor<InT> o({1, 2, 4, 8});
|
||
|
||
q.fillWithValue(safeTestTypeCast<InT>(0.1f));
|
||
k.fillWithValue(safeTestTypeCast<InT>(0.1f));
|
||
|
||
const float vVal = 0.5f;
|
||
v.fillWithValue(safeTestTypeCast<InT>(vVal));
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// With uniform scores, softmax gives equal weights 1/Skv for each kv token.
|
||
// With uniform V = vVal, O = sum_skv(1/Skv * vVal) = vVal.
|
||
// Use a generous tolerance to accommodate half/bfloat16 precision limits.
|
||
const float tol = 0.05f;
|
||
EXPECT_NEAR(static_cast<float>(o.getHostValue(0, 0, 0, 0)), vVal, tol);
|
||
EXPECT_NEAR(static_cast<float>(o.getHostValue(0, 1, 3, 7)), vVal, tol);
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// LSE (Log-Sum-Exp) Output Tests
|
||
// ---------------------------------------------------------------------------
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseOutputMatchesFormula)
|
||
{
|
||
// Small controlled input: [B=1, H=1, Sq=2, Skv=2, D=2, Dv=2]
|
||
// Use one-hot Q/K for known scores, then verify LSE = maxVal + log(sumExp)
|
||
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 1, 2, 2});
|
||
Tensor<float> lse({1, 1, 2, 1});
|
||
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
|
||
// One-hot pattern: Q[0,0,0,:] = [1,0], Q[0,0,1,:] = [0,1]
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
q.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
// K[0,0,0,:] = [1,0], K[0,0,1,:] = [0,1]
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 1, 1);
|
||
|
||
v.setHostValue(1.0, 0, 0, 0, 0);
|
||
v.setHostValue(2.0, 0, 0, 0, 1);
|
||
v.setHostValue(3.0, 0, 0, 1, 0);
|
||
v.setHostValue(4.0, 0, 0, 1, 1);
|
||
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lse);
|
||
|
||
// Manually compute expected LSE
|
||
const float scale = 1.0f / std::sqrt(2.0f);
|
||
|
||
// For sq=0: dot products are [1,0] → scores = [scale, 0]
|
||
const float maxVal0 = scale;
|
||
const float sumExp0 = std::exp(scale - maxVal0) + std::exp(0.0f - maxVal0);
|
||
const float expectedLse0 = maxVal0 + std::log(sumExp0);
|
||
|
||
EXPECT_NEAR(lse.getHostValue(0, 0, 0, 0), expectedLse0, 1e-5f);
|
||
|
||
// For sq=1: dot products are [0,1] → scores = [0, scale]
|
||
const float maxVal1 = scale;
|
||
const float sumExp1 = std::exp(0.0f - maxVal1) + std::exp(scale - maxVal1);
|
||
const float expectedLse1 = maxVal1 + std::log(sumExp1);
|
||
|
||
EXPECT_NEAR(lse.getHostValue(0, 0, 1, 0), expectedLse1, 1e-5f);
|
||
}
|
||
|
||
TYPED_TEST(CpuFpReferenceSdpaFwd, LseAlwaysFloatType)
|
||
{
|
||
// Verify: LSE is always float, even when inputs are half/bfloat16
|
||
using InT = typename TypeParam::First;
|
||
|
||
Tensor<InT> q({1, 2, 4, 8});
|
||
Tensor<InT> k({1, 2, 4, 8});
|
||
Tensor<InT> v({1, 2, 4, 8});
|
||
Tensor<InT> o({1, 2, 4, 8});
|
||
Tensor<float> lse({1, 2, 4, 1});
|
||
|
||
q.fillWithRandomValues(safeTestTypeCast<InT>(-1.0f), safeTestTypeCast<InT>(1.0f), 100);
|
||
k.fillWithRandomValues(safeTestTypeCast<InT>(-1.0f), safeTestTypeCast<InT>(1.0f), 101);
|
||
v.fillWithRandomValues(safeTestTypeCast<InT>(-1.0f), safeTestTypeCast<InT>(1.0f), 102);
|
||
|
||
const hipdnn_data_sdk::utilities::TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lse);
|
||
|
||
// LSE values should be reasonable (not NaN, not inf for non-masked rows)
|
||
for(int b = 0; b < 1; ++b)
|
||
{
|
||
for(int h = 0; h < 2; ++h)
|
||
{
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
const float lseVal = lse.getHostValue(b, h, sq, 0);
|
||
EXPECT_FALSE(std::isnan(lseVal)) << "NaN at [" << b << "," << h << "," << sq << "]";
|
||
// LSE typically in range [-10, 10] for random inputs with default scale
|
||
EXPECT_GT(lseVal, -20.0f)
|
||
<< "Unexpectedly small LSE at [" << b << "," << h << "," << sq << "]";
|
||
EXPECT_LT(lseVal, 20.0f)
|
||
<< "Unexpectedly large LSE at [" << b << "," << h << "," << sq << "]";
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWithMultipleBatchHeads)
|
||
{
|
||
// Verify: LSE computed correctly for B > 1, H > 1
|
||
|
||
Tensor<double> q({2, 4, 16, 32});
|
||
Tensor<double> k({2, 4, 16, 32});
|
||
Tensor<double> v({2, 4, 16, 32});
|
||
Tensor<double> o({2, 4, 16, 32});
|
||
Tensor<float> lse({2, 4, 16, 1});
|
||
|
||
q.fillWithRandomValues(-1.0, 1.0, 200);
|
||
k.fillWithRandomValues(-1.0, 1.0, 201);
|
||
v.fillWithRandomValues(-1.0, 1.0, 202);
|
||
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lse);
|
||
|
||
// Verify all LSE values are reasonable (no NaN, finite)
|
||
for(int b = 0; b < 2; ++b)
|
||
{
|
||
for(int h = 0; h < 4; ++h)
|
||
{
|
||
for(int sq = 0; sq < 16; ++sq)
|
||
{
|
||
const float lseVal = lse.getHostValue(b, h, sq, 0);
|
||
EXPECT_TRUE(std::isfinite(lseVal))
|
||
<< "LSE not finite at [" << b << "," << h << "," << sq << "]";
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWithCausalMask)
|
||
{
|
||
// [B=1, H=1, Sq=4, Skv=4]
|
||
// Causal mask: sq=0 sees 1 position, sq=1 sees 2, sq=2 sees 3, sq=3 sees 4
|
||
|
||
Tensor<double> q({1, 1, 4, 8});
|
||
Tensor<double> k({1, 1, 4, 8});
|
||
Tensor<double> v({1, 1, 4, 8});
|
||
Tensor<double> o({1, 1, 4, 8});
|
||
Tensor<float> lse({1, 1, 4, 1});
|
||
|
||
q.fillWithValue(1.0);
|
||
k.fillWithValue(1.0);
|
||
v.fillWithValue(1.0);
|
||
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, true, &lse);
|
||
|
||
// With uniform Q/K and causal mask:
|
||
// sq=0: 1 valid position → smaller sumExp → smaller LSE
|
||
// sq=3: 4 valid positions → larger sumExp → larger LSE
|
||
// LSE should increase monotonically as sq increases
|
||
|
||
const float lse0 = lse.getHostValue(0, 0, 0, 0);
|
||
const float lse1 = lse.getHostValue(0, 0, 1, 0);
|
||
const float lse2 = lse.getHostValue(0, 0, 2, 0);
|
||
const float lse3 = lse.getHostValue(0, 0, 3, 0);
|
||
|
||
EXPECT_LT(lse0, lse1) << "LSE should increase with more unmasked positions";
|
||
EXPECT_LT(lse1, lse2);
|
||
EXPECT_LT(lse2, lse3);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWithFullyMaskedRow)
|
||
{
|
||
// Edge case: All scores masked to -inf
|
||
// LSE should be -inf (mathematically correct: log(0) = -inf)
|
||
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 1, 2, 2});
|
||
Tensor<float> lse({1, 1, 2, 1});
|
||
|
||
q.fillWithValue(1.0);
|
||
k.fillWithValue(1.0);
|
||
v.fillWithValue(1.0);
|
||
|
||
// Create additive mask that blocks ALL positions for sq=0
|
||
Tensor<float> mask({2, 2});
|
||
mask.fillWithValue(-std::numeric_limits<float>::infinity());
|
||
|
||
// But allow sq=1 to see all positions
|
||
mask.setHostValue(0.0f, 1, 0);
|
||
mask.setHostValue(0.0f, 1, 1);
|
||
|
||
const hipdnn_data_sdk::utilities::TensorBase<float>* maskPtr = &mask;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, maskPtr, false, &lse);
|
||
|
||
// sq=0: all masked → LSE = -inf (or NaN due to -inf - (-inf) in exp)
|
||
const float lse0 = lse.getHostValue(0, 0, 0, 0);
|
||
// When all scores are -inf: maxVal = -inf, exp(-inf - (-inf)) = exp(NaN) = NaN
|
||
// So LSE may be NaN rather than -inf. Both indicate "no valid attention weights"
|
||
EXPECT_TRUE((std::isinf(lse0) && lse0 < 0) || std::isnan(lse0))
|
||
<< "LSE should be -inf or NaN for fully masked row, got: " << lse0;
|
||
|
||
// sq=1: normal → LSE finite
|
||
const float lse1 = lse.getHostValue(0, 0, 1, 0);
|
||
EXPECT_TRUE(std::isfinite(lse1)) << "LSE should be finite for normal row";
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWithAdditiveMask)
|
||
{
|
||
// Verify: LSE accounts for additive mask contributions
|
||
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o1({1, 1, 2, 2});
|
||
Tensor<double> o2({1, 1, 2, 2});
|
||
Tensor<float> lse1({1, 1, 2, 1});
|
||
Tensor<float> lse2({1, 1, 2, 1});
|
||
|
||
// One-hot Q/K for predictable scores
|
||
q.fillWithValue(0.0);
|
||
k.fillWithValue(0.0);
|
||
q.setHostValue(1.0, 0, 0, 0, 0);
|
||
q.setHostValue(1.0, 0, 0, 1, 1);
|
||
k.setHostValue(1.0, 0, 0, 0, 0);
|
||
k.setHostValue(1.0, 0, 0, 1, 1);
|
||
v.fillWithValue(1.0);
|
||
|
||
// Forward without mask
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o1, std::nullopt, noMask, false, &lse1);
|
||
|
||
// Forward with mask that suppresses one position
|
||
Tensor<float> mask({2, 2});
|
||
mask.fillWithValue(0.0f);
|
||
mask.setHostValue(-1e4f, 0, 1); // Block skv=1 for sq=0
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o2, std::nullopt, &mask, false, &lse2);
|
||
|
||
// LSE should differ between masked and unmasked cases
|
||
const float lseNomask = lse1.getHostValue(0, 0, 0, 0);
|
||
const float lseWithmask = lse2.getHostValue(0, 0, 0, 0);
|
||
|
||
EXPECT_NE(lseNomask, lseWithmask) << "LSE should change with additive mask";
|
||
|
||
// With mask blocking skv=1, LSE should be lower (fewer effective positions)
|
||
EXPECT_LT(lseWithmask, lseNomask) << "LSE should decrease when positions are masked";
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWrongRank)
|
||
{
|
||
Tensor<double> q({1, 1, 2, 2});
|
||
Tensor<double> k({1, 1, 2, 2});
|
||
Tensor<double> v({1, 1, 2, 2});
|
||
Tensor<double> o({1, 1, 2, 2});
|
||
|
||
// Wrong rank: rank-3 instead of rank-4
|
||
Tensor<float> lseWrong({1, 1, 2});
|
||
|
||
EXPECT_THROW(
|
||
{
|
||
q.fillWithValue(1.0);
|
||
k.fillWithValue(1.0);
|
||
v.fillWithValue(1.0);
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lseWrong);
|
||
},
|
||
std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWrongShape)
|
||
{
|
||
Tensor<double> q({2, 4, 16, 32});
|
||
Tensor<double> k({2, 4, 16, 32});
|
||
Tensor<double> v({2, 4, 16, 32});
|
||
Tensor<double> o({2, 4, 16, 32});
|
||
|
||
// Wrong shape: [2, 4, 8, 1] instead of [2, 4, 16, 1]
|
||
Tensor<float> lseWrong({2, 4, 8, 1});
|
||
|
||
EXPECT_THROW(
|
||
{
|
||
q.fillWithRandomValues(-1.0, 1.0, 42);
|
||
k.fillWithRandomValues(-1.0, 1.0, 43);
|
||
v.fillWithRandomValues(-1.0, 1.0, 44);
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lseWrong);
|
||
},
|
||
std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseMismatchedBatch)
|
||
{
|
||
Tensor<double> q({2, 1, 4, 8});
|
||
Tensor<double> k({2, 1, 4, 8});
|
||
Tensor<double> v({2, 1, 4, 8});
|
||
Tensor<double> o({2, 1, 4, 8});
|
||
|
||
// Wrong batch: [1, 1, 4, 1] instead of [2, 1, 4, 1]
|
||
Tensor<float> lseWrong({1, 1, 4, 1});
|
||
|
||
EXPECT_THROW(
|
||
{
|
||
q.fillWithRandomValues(-1.0, 1.0, 50);
|
||
k.fillWithRandomValues(-1.0, 1.0, 51);
|
||
v.fillWithRandomValues(-1.0, 1.0, 52);
|
||
const hipdnn_data_sdk::utilities::TensorBase<double>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lseWrong);
|
||
},
|
||
std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaFp64, LseWithGqa)
|
||
{
|
||
// GQA: H=8, HkV=2 (headsPerKvHead = 4)
|
||
// Verify: LSE computed correctly despite K/V head sharing
|
||
|
||
Tensor<double> q({1, 8, 16, 32});
|
||
Tensor<double> k({1, 2, 16, 32});
|
||
Tensor<double> v({1, 2, 16, 32});
|
||
Tensor<double> o({1, 8, 16, 32});
|
||
Tensor<float> lse({1, 8, 16, 1});
|
||
|
||
q.fillWithRandomValues(-1.0, 1.0, 300);
|
||
k.fillWithRandomValues(-1.0, 1.0, 301);
|
||
v.fillWithRandomValues(-1.0, 1.0, 302);
|
||
|
||
const hipdnn_data_sdk::utilities::TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lse);
|
||
|
||
// All LSE values should be reasonable
|
||
for(int h = 0; h < 8; ++h)
|
||
{
|
||
for(int sq = 0; sq < 16; ++sq)
|
||
{
|
||
const float lseVal = lse.getHostValue(0, h, sq, 0);
|
||
EXPECT_TRUE(std::isfinite(lseVal)) << "LSE not finite at h=" << h << ", sq=" << sq;
|
||
EXPECT_GT(lseVal, -20.0f);
|
||
EXPECT_LT(lseVal, 20.0f);
|
||
}
|
||
}
|
||
|
||
// Different Q heads using same KV head should have different LSE
|
||
// (because Q differs, even though K/V are shared)
|
||
const float lseH0 = lse.getHostValue(0, 0, 0, 0);
|
||
const float lseH1 = lse.getHostValue(0, 1, 0, 0);
|
||
|
||
// With random Q, LSE should almost certainly differ
|
||
EXPECT_NE(lseH0, lseH1) << "LSE should differ for different Q heads even with shared KV";
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Backward pass tests
|
||
//
|
||
// These tests use Fp32 (unlike forward tests which use Fp64) because the GPU
|
||
// backward kernels (AITER/Flash Attention) only support fp16/bf16 input/output
|
||
// with fp32 internal compute precision. The CPU reference backward tests use
|
||
// Fp32 to match the actual kernel compute precision.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardSanity)
|
||
{
|
||
// Hand-verified backward pass with small deterministic inputs.
|
||
// [B=1, H=1, Sq=2, Skv=2, D=2, Dv=2]
|
||
//
|
||
// Q = [[1,0],[0,1]] K = [[1,0],[0,1]] V = [[1,2],[3,4]] dO = [[1,1],[1,1]]
|
||
// scale = 1/sqrt(2)
|
||
//
|
||
// By symmetry of Q and K (identity matrices), the softmax probabilities
|
||
// for sq=0 are [pH, pL] and for sq=1 are [pL, pH], where:
|
||
// pH = exp(scale) / (exp(scale) + 1) (probability of the matching key)
|
||
// pL = 1 / (exp(scale) + 1) (probability of the non-matching key)
|
||
//
|
||
// With dO = ones, dP[sq,skv] = sum(V[skv,:]) = [3, 7].
|
||
// dS[sq,skv] = P[sq,skv] * (dP[sq,skv] - D[sq]) simplifies to:
|
||
// dS[*,matching] = -4*pH*pL, dS[*,non-matching] = +4*pH*pL
|
||
//
|
||
// Expected gradients (c = 4*pH*pL*scale):
|
||
// dQ = [[-c, c], [-c, c]]
|
||
// dK = [[-c, -c], [c, c]]
|
||
// dV = [[1, 1], [1, 1]]
|
||
Tensor<float> q({1, 1, 2, 2});
|
||
Tensor<float> k({1, 1, 2, 2});
|
||
Tensor<float> v({1, 1, 2, 2});
|
||
Tensor<float> o({1, 1, 2, 2});
|
||
Tensor<float> dO({1, 1, 2, 2});
|
||
Tensor<float> dQ({1, 1, 2, 2});
|
||
Tensor<float> dK({1, 1, 2, 2});
|
||
Tensor<float> dV({1, 1, 2, 2});
|
||
|
||
q.fillWithValue(0.0f);
|
||
k.fillWithValue(0.0f);
|
||
v.fillWithValue(0.0f);
|
||
|
||
q.setHostValue(1.0f, 0, 0, 0, 0);
|
||
q.setHostValue(1.0f, 0, 0, 1, 1);
|
||
k.setHostValue(1.0f, 0, 0, 0, 0);
|
||
k.setHostValue(1.0f, 0, 0, 1, 1);
|
||
v.setHostValue(1.0f, 0, 0, 0, 0);
|
||
v.setHostValue(2.0f, 0, 0, 0, 1);
|
||
v.setHostValue(3.0f, 0, 0, 1, 0);
|
||
v.setHostValue(4.0f, 0, 0, 1, 1);
|
||
|
||
dO.fillWithValue(1.0f);
|
||
|
||
// Forward pass
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Backward pass
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
// Compute expected values
|
||
const float scale = 1.0f / std::sqrt(2.0f);
|
||
const float eS = std::exp(scale);
|
||
const float pH = eS / (eS + 1.0f);
|
||
const float pL = 1.0f / (eS + 1.0f);
|
||
const float c = 4.0f * pH * pL * scale;
|
||
const float tol = 1e-5f;
|
||
|
||
// dQ = [[-c, c], [-c, c]]
|
||
EXPECT_NEAR(dQ.getHostValue(0, 0, 0, 0), -c, tol);
|
||
EXPECT_NEAR(dQ.getHostValue(0, 0, 0, 1), c, tol);
|
||
EXPECT_NEAR(dQ.getHostValue(0, 0, 1, 0), -c, tol);
|
||
EXPECT_NEAR(dQ.getHostValue(0, 0, 1, 1), c, tol);
|
||
|
||
// dK = [[-c, -c], [c, c]]
|
||
EXPECT_NEAR(dK.getHostValue(0, 0, 0, 0), -c, tol);
|
||
EXPECT_NEAR(dK.getHostValue(0, 0, 0, 1), -c, tol);
|
||
EXPECT_NEAR(dK.getHostValue(0, 0, 1, 0), c, tol);
|
||
EXPECT_NEAR(dK.getHostValue(0, 0, 1, 1), c, tol);
|
||
|
||
// dV = [[1, 1], [1, 1]] (P columns sum to 1 by softmax symmetry)
|
||
EXPECT_NEAR(dV.getHostValue(0, 0, 0, 0), 1.0f, tol);
|
||
EXPECT_NEAR(dV.getHostValue(0, 0, 0, 1), 1.0f, tol);
|
||
EXPECT_NEAR(dV.getHostValue(0, 0, 1, 0), 1.0f, tol);
|
||
EXPECT_NEAR(dV.getHostValue(0, 0, 1, 1), 1.0f, tol);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardMHA)
|
||
{
|
||
// Test Multi-Head Attention (H_q == H_kv)
|
||
// [B=1, H=2, Sq=4, Skv=4, D=8, Dv=8]
|
||
Tensor<float> q({1, 2, 4, 8});
|
||
Tensor<float> k({1, 2, 4, 8});
|
||
Tensor<float> v({1, 2, 4, 8});
|
||
Tensor<float> o({1, 2, 4, 8});
|
||
Tensor<float> dO({1, 2, 4, 8});
|
||
Tensor<float> dQ({1, 2, 4, 8});
|
||
Tensor<float> dK({1, 2, 4, 8});
|
||
Tensor<float> dV({1, 2, 4, 8});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 100);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 101);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 102);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 103);
|
||
|
||
// Forward pass
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Backward pass
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
// Check gradients are finite
|
||
for(int h = 0; h < 2; ++h)
|
||
{
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 8; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dQ.getHostValue(0, h, sq, d)))
|
||
<< "dQ not finite at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
EXPECT_TRUE(std::isfinite(dK.getHostValue(0, h, sq, d)))
|
||
<< "dK not finite at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
EXPECT_TRUE(std::isfinite(dV.getHostValue(0, h, sq, d)))
|
||
<< "dV not finite at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardGQA)
|
||
{
|
||
// Test Grouped Query Attention (H_q = 4, H_kv = 1)
|
||
// [B=1, H_q=4, H_kv=1, Sq=4, Skv=4, D=8, Dv=8]
|
||
Tensor<float> q({1, 4, 4, 8});
|
||
Tensor<float> k({1, 1, 4, 8});
|
||
Tensor<float> v({1, 1, 4, 8});
|
||
Tensor<float> o({1, 4, 4, 8});
|
||
Tensor<float> dO({1, 4, 4, 8});
|
||
Tensor<float> dQ({1, 4, 4, 8});
|
||
Tensor<float> dK({1, 1, 4, 8});
|
||
Tensor<float> dV({1, 1, 4, 8});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 200);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 201);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 202);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 203);
|
||
|
||
// Forward pass
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Backward pass (should handle GQA accumulation correctly)
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
// Check gradients are finite
|
||
for(int h = 0; h < 4; ++h)
|
||
{
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 8; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dQ.getHostValue(0, h, sq, d)))
|
||
<< "dQ not finite at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 8; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dK.getHostValue(0, 0, sq, d)))
|
||
<< "dK not finite at sq=" << sq << ", d=" << d;
|
||
EXPECT_TRUE(std::isfinite(dV.getHostValue(0, 0, sq, d)))
|
||
<< "dV not finite at sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
|
||
// For GQA, dK and dV should accumulate contributions from multiple Q heads
|
||
// Verify that gradients are non-zero (as they should accumulate from 4 Q heads)
|
||
float dkSum = 0.0f;
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 8; ++d)
|
||
{
|
||
dkSum += std::abs(dK.getHostValue(0, 0, sq, d));
|
||
}
|
||
}
|
||
EXPECT_GT(dkSum, 0.0f) << "dK should have non-zero gradients from GQA accumulation";
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardCausalMask)
|
||
{
|
||
// Test backward pass with causal masking
|
||
// [B=1, H=1, Sq=4, Skv=4, D=4, Dv=4]
|
||
Tensor<float> q({1, 1, 4, 4});
|
||
Tensor<float> k({1, 1, 4, 4});
|
||
Tensor<float> v({1, 1, 4, 4});
|
||
Tensor<float> o({1, 1, 4, 4});
|
||
Tensor<float> dO({1, 1, 4, 4});
|
||
Tensor<float> dQ({1, 1, 4, 4});
|
||
Tensor<float> dK({1, 1, 4, 4});
|
||
Tensor<float> dV({1, 1, 4, 4});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 300);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 301);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 302);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 303);
|
||
|
||
const bool causalMask = true;
|
||
const hipdnn_data_sdk::utilities::TensorBase<float>* noMask = nullptr;
|
||
|
||
// Forward pass with causal mask
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, causalMask);
|
||
|
||
// Backward pass with causal mask
|
||
CpuFpReferenceSdpa::backward(
|
||
q, k, v, o, dO, dQ, dK, dV, std::nullopt, nullptr, noMask, causalMask);
|
||
|
||
// Check gradients are finite
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dQ.getHostValue(0, 0, sq, d)))
|
||
<< "dQ not finite at sq=" << sq << ", d=" << d;
|
||
EXPECT_TRUE(std::isfinite(dK.getHostValue(0, 0, sq, d)))
|
||
<< "dK not finite at sq=" << sq << ", d=" << d;
|
||
EXPECT_TRUE(std::isfinite(dV.getHostValue(0, 0, sq, d)))
|
||
<< "dV not finite at sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Compares LSE-based backward path against full softmax recomputation path.
|
||
// Both paths must use float tensors (matching ComputeDataType = float default)
|
||
// so that forward and backward compute at the same precision.
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardWithLSE)
|
||
{
|
||
// [B=1, H=2, Sq=4, Skv=4, D=8, Dv=8]
|
||
Tensor<float> q({1, 2, 4, 8});
|
||
Tensor<float> k({1, 2, 4, 8});
|
||
Tensor<float> v({1, 2, 4, 8});
|
||
Tensor<float> o({1, 2, 4, 8});
|
||
Tensor<float> lse({1, 2, 4, 1});
|
||
Tensor<float> dO({1, 2, 4, 8});
|
||
Tensor<float> dQWithLse({1, 2, 4, 8});
|
||
Tensor<float> dKWithLse({1, 2, 4, 8});
|
||
Tensor<float> dVWithLse({1, 2, 4, 8});
|
||
Tensor<float> dQNoLse({1, 2, 4, 8});
|
||
Tensor<float> dKNoLse({1, 2, 4, 8});
|
||
Tensor<float> dVNoLse({1, 2, 4, 8});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 400);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 401);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 402);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 403);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
|
||
// Forward pass with LSE
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, false, &lse);
|
||
|
||
// Backward pass WITH LSE (efficient)
|
||
CpuFpReferenceSdpa::backward(
|
||
q, k, v, o, dO, dQWithLse, dKWithLse, dVWithLse, std::nullopt, &lse);
|
||
|
||
// Backward pass WITHOUT LSE (recompute from scratch)
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQNoLse, dKNoLse, dVNoLse);
|
||
|
||
// Both paths compute in float, so results should be very close.
|
||
// LSE path uses exp(score - lse) while no-LSE path uses exp(score - max) / sumExp.
|
||
// These differ by FP operation ordering. If this tolerance causes flaky failures,
|
||
// it can be relaxed to 1e-4f.
|
||
const float tol = 1e-5f;
|
||
for(int h = 0; h < 2; ++h)
|
||
{
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 8; ++d)
|
||
{
|
||
EXPECT_NEAR(
|
||
dQWithLse.getHostValue(0, h, sq, d), dQNoLse.getHostValue(0, h, sq, d), tol)
|
||
<< "dQ mismatch at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
EXPECT_NEAR(
|
||
dKWithLse.getHostValue(0, h, sq, d), dKNoLse.getHostValue(0, h, sq, d), tol)
|
||
<< "dK mismatch at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
EXPECT_NEAR(
|
||
dVWithLse.getHostValue(0, h, sq, d), dVNoLse.getHostValue(0, h, sq, d), tol)
|
||
<< "dV mismatch at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Finite-difference gradient checking
|
||
// ---------------------------------------------------------------------------
|
||
//
|
||
// Validates backward() by comparing its analytical gradients against numerical
|
||
// gradients computed via central finite differences on the forward pass:
|
||
//
|
||
// dL/dX[i] ≈ (L(X + ε·e_i) − L(X − ε·e_i)) / (2ε)
|
||
//
|
||
// where L = Σ(dO ⊙ O) is a scalar loss and O = forward(Q, K, V).
|
||
// This is the gold-standard test for gradient correctness.
|
||
|
||
using hipdnn_test_sdk::detail::compareGradients;
|
||
using hipdnn_test_sdk::detail::computeDotProductLoss;
|
||
using hipdnn_test_sdk::detail::numericalGradient;
|
||
|
||
static void checkGradient(Tensor<float>& input,
|
||
const Tensor<float>& analyticalGrad,
|
||
const std::vector<int64_t>& shape,
|
||
const std::function<double()>& fwdLoss,
|
||
const std::string& label,
|
||
double eps = 1e-3,
|
||
double relTol = 1e-2,
|
||
double absTol = 1e-4)
|
||
{
|
||
Tensor<float> numericalGrad(shape);
|
||
numericalGradient(input, numericalGrad, eps, fwdLoss);
|
||
auto result = compareGradients(analyticalGrad, numericalGrad, relTol, absTol);
|
||
EXPECT_EQ(result.failCount, 0) << label << " grad check failed: maxRelErr=" << result.maxRelErr
|
||
<< ", maxAbsErr=" << result.maxAbsErr;
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, MHA)
|
||
{
|
||
// MHA: H_q == H_kv, no masks
|
||
// [B=1, H=2, Sq=3, Skv=3, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 2;
|
||
constexpr int64_t SEQ_Q = 3;
|
||
constexpr int64_t SEQ_KV = 3;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 1000);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 1001);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 1002);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 1003);
|
||
|
||
// Analytical backward
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, GQA)
|
||
{
|
||
// GQA: H_q=4, H_kv=1 — tests KV gradient accumulation across Q heads
|
||
// [B=1, H_q=4, H_kv=1, Sq=3, Skv=3, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS_Q = 4;
|
||
constexpr int64_t HEADS_KV = 1;
|
||
constexpr int64_t SEQ_Q = 3;
|
||
constexpr int64_t SEQ_KV = 3;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 2000);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 2001);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 2002);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 2003);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS_Q, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ");
|
||
checkGradient(k, dK, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM}, fwdLoss, "dK");
|
||
checkGradient(v, dV, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, CausalMask)
|
||
{
|
||
// MHA with causal masking
|
||
// [B=1, H=1, Sq=4, Skv=4, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 4;
|
||
constexpr int64_t SEQ_KV = 4;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 3000);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 3001);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 3002);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 3003);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, /*causalMask=*/true);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(
|
||
q, k, v, o, dO, dQ, dK, dV, std::nullopt, nullptr, noMask, /*causalMask=*/true);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, /*causalMask=*/true);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (causal)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (causal)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (causal)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, CustomScale)
|
||
{
|
||
// MHA with a non-default attention scale
|
||
// [B=1, H=1, Sq=3, Skv=3, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 3;
|
||
constexpr int64_t SEQ_KV = 3;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 4000);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 4001);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 4002);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 4003);
|
||
|
||
const auto customScale = std::optional<float>{0.3f};
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, customScale);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV, customScale);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, customScale);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (scale=0.3)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (scale=0.3)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (scale=0.3)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, AdditiveMask)
|
||
{
|
||
// MHA with an additive attention mask
|
||
// [B=1, H=1, Sq=3, Skv=3, D=4, Dv=4], mask [Sq, Skv]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 3;
|
||
constexpr int64_t SEQ_KV = 3;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> mask({SEQ_Q, SEQ_KV});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 5000);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 5001);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 5002);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 5003);
|
||
|
||
// Partially suppress some positions
|
||
mask.fillWithValue(0.0f);
|
||
mask.setHostValue(-1e4f, 0, 2); // Block sq=0 attending to skv=2
|
||
mask.setHostValue(-1e4f, 1, 0); // Block sq=1 attending to skv=0
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, &mask);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV, std::nullopt, nullptr, &mask, false);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, &mask);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (additive mask)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (additive mask)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (additive mask)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, GQACausalMask)
|
||
{
|
||
// Combined: GQA + causal mask
|
||
// [B=1, H_q=4, H_kv=2, Sq=4, Skv=4, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS_Q = 4;
|
||
constexpr int64_t HEADS_KV = 2;
|
||
constexpr int64_t SEQ_Q = 4;
|
||
constexpr int64_t SEQ_KV = 4;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 6000);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 6001);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 6002);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 6003);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, /*causalMask=*/true);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(
|
||
q, k, v, o, dO, dQ, dK, dV, std::nullopt, nullptr, noMask, /*causalMask=*/true);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, /*causalMask=*/true);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS_Q, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (GQA+causal)");
|
||
checkGradient(k, dK, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (GQA+causal)");
|
||
checkGradient(v, dV, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (GQA+causal)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdBf16, BackwardBasic)
|
||
{
|
||
// Test backward pass with BFloat16 data type
|
||
// [B=1, H=1, Sq=4, Skv=4, D=8, Dv=8]
|
||
Tensor<bfloat16> q({1, 1, 4, 8});
|
||
Tensor<bfloat16> k({1, 1, 4, 8});
|
||
Tensor<bfloat16> v({1, 1, 4, 8});
|
||
Tensor<bfloat16> o({1, 1, 4, 8});
|
||
Tensor<bfloat16> dO({1, 1, 4, 8});
|
||
Tensor<bfloat16> dQ({1, 1, 4, 8});
|
||
Tensor<bfloat16> dK({1, 1, 4, 8});
|
||
Tensor<bfloat16> dV({1, 1, 4, 8});
|
||
|
||
q.fillWithRandomValues(bfloat16(-1.0f), bfloat16(1.0f), 500);
|
||
k.fillWithRandomValues(bfloat16(-1.0f), bfloat16(1.0f), 501);
|
||
v.fillWithRandomValues(bfloat16(-1.0f), bfloat16(1.0f), 502);
|
||
dO.fillWithRandomValues(bfloat16(-1.0f), bfloat16(1.0f), 503);
|
||
|
||
// Forward pass
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Backward pass
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
// Check gradients are finite (BFloat16 has limited precision)
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 8; ++d)
|
||
{
|
||
const float dqVal = static_cast<float>(dQ.getHostValue(0, 0, sq, d));
|
||
EXPECT_TRUE(std::isfinite(dqVal)) << "dQ not finite at sq=" << sq << ", d=" << d;
|
||
const float dkVal = static_cast<float>(dK.getHostValue(0, 0, sq, d));
|
||
EXPECT_TRUE(std::isfinite(dkVal)) << "dK not finite at sq=" << sq << ", d=" << d;
|
||
const float dvVal = static_cast<float>(dV.getHostValue(0, 0, sq, d));
|
||
EXPECT_TRUE(std::isfinite(dvVal)) << "dV not finite at sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardGqaDifferentKVHeads)
|
||
{
|
||
// Test backward with independent K and V head counts: H_q=4, H_k=2, H_v=1
|
||
// Mirrors the forward GqaDifferentKVHeads test.
|
||
// [B=1, H_q=4, H_k=2, H_v=1, Sq=3, Skv=3, D=4, Dv=4]
|
||
Tensor<float> q({1, 4, 3, 4});
|
||
Tensor<float> k({1, 2, 3, 4});
|
||
Tensor<float> v({1, 1, 3, 4});
|
||
Tensor<float> o({1, 4, 3, 4});
|
||
Tensor<float> dO({1, 4, 3, 4});
|
||
Tensor<float> dQ({1, 4, 3, 4});
|
||
Tensor<float> dK({1, 2, 3, 4});
|
||
Tensor<float> dV({1, 1, 3, 4});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 600);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 601);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 602);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 603);
|
||
|
||
// Forward pass
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
// Backward pass with different K and V head counts
|
||
CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV);
|
||
|
||
// Check dQ gradients are finite (4 Q heads)
|
||
for(int h = 0; h < 4; ++h)
|
||
{
|
||
for(int sq = 0; sq < 3; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dQ.getHostValue(0, h, sq, d)))
|
||
<< "dQ not finite at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Check dK gradients are finite (2 K heads)
|
||
for(int h = 0; h < 2; ++h)
|
||
{
|
||
for(int sq = 0; sq < 3; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dK.getHostValue(0, h, sq, d)))
|
||
<< "dK not finite at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Check dV gradients are finite (1 V head)
|
||
for(int sq = 0; sq < 3; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
EXPECT_TRUE(std::isfinite(dV.getHostValue(0, 0, sq, d)))
|
||
<< "dV not finite at sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
|
||
// dK should accumulate from 2 Q heads per K head
|
||
float dkSum = 0.0f;
|
||
for(int h = 0; h < 2; ++h)
|
||
{
|
||
for(int sq = 0; sq < 3; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
dkSum += std::abs(dK.getHostValue(0, h, sq, d));
|
||
}
|
||
}
|
||
}
|
||
EXPECT_GT(dkSum, 0.0f) << "dK should have non-zero gradients from GQA accumulation";
|
||
|
||
// dV should accumulate from all 4 Q heads into the single V head
|
||
float dvSum = 0.0f;
|
||
for(int sq = 0; sq < 3; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
dvSum += std::abs(dV.getHostValue(0, 0, sq, d));
|
||
}
|
||
}
|
||
EXPECT_GT(dvSum, 0.0f) << "dV should have non-zero gradients from GQA accumulation";
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Backward validation tests (error cases)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardWrongRankQ)
|
||
{
|
||
const Tensor<float> q({1, 2, 4}); // rank-3 instead of rank-4
|
||
const Tensor<float> k({1, 1, 4, 8});
|
||
const Tensor<float> v({1, 1, 4, 8});
|
||
const Tensor<float> o({1, 1, 4, 8});
|
||
const Tensor<float> dO({1, 1, 4, 8});
|
||
Tensor<float> dQ({1, 2, 4});
|
||
Tensor<float> dK({1, 1, 4, 8});
|
||
Tensor<float> dV({1, 1, 4, 8});
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV), std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardWrongRankDK)
|
||
{
|
||
const Tensor<float> q({1, 1, 4, 8});
|
||
const Tensor<float> k({1, 1, 4, 8});
|
||
const Tensor<float> v({1, 1, 4, 8});
|
||
const Tensor<float> o({1, 1, 4, 8});
|
||
const Tensor<float> dO({1, 1, 4, 8});
|
||
Tensor<float> dQ({1, 1, 4, 8});
|
||
Tensor<float> dK({1, 4, 8}); // rank-3 instead of rank-4
|
||
Tensor<float> dV({1, 1, 4, 8});
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV), std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardBatchMismatch)
|
||
{
|
||
const Tensor<float> q({2, 1, 4, 8});
|
||
const Tensor<float> k({2, 1, 4, 8});
|
||
const Tensor<float> v({2, 1, 4, 8});
|
||
const Tensor<float> o({2, 1, 4, 8});
|
||
const Tensor<float> dO({2, 1, 4, 8});
|
||
Tensor<float> dQ({2, 1, 4, 8});
|
||
Tensor<float> dK({1, 1, 4, 8}); // batch=1 instead of batch=2
|
||
Tensor<float> dV({2, 1, 4, 8});
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV), std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardHeadDimMismatch)
|
||
{
|
||
const Tensor<float> q({1, 1, 4, 8});
|
||
const Tensor<float> k({1, 1, 4, 16}); // D=16 instead of D=8
|
||
const Tensor<float> v({1, 1, 4, 8});
|
||
const Tensor<float> o({1, 1, 4, 8});
|
||
const Tensor<float> dO({1, 1, 4, 8});
|
||
Tensor<float> dQ({1, 1, 4, 8});
|
||
Tensor<float> dK({1, 1, 4, 16});
|
||
Tensor<float> dV({1, 1, 4, 8});
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV), std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardGqaNotDivisibleK)
|
||
{
|
||
// H_q=3 is not divisible by H_k=2
|
||
const Tensor<float> q({1, 3, 4, 8});
|
||
const Tensor<float> k({1, 2, 4, 8});
|
||
const Tensor<float> v({1, 1, 4, 8});
|
||
const Tensor<float> o({1, 3, 4, 8});
|
||
const Tensor<float> dO({1, 3, 4, 8});
|
||
Tensor<float> dQ({1, 3, 4, 8});
|
||
Tensor<float> dK({1, 2, 4, 8});
|
||
Tensor<float> dV({1, 1, 4, 8});
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV), std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardGqaNotDivisibleV)
|
||
{
|
||
// H_q=4 is not divisible by H_v=3
|
||
const Tensor<float> q({1, 4, 4, 8});
|
||
const Tensor<float> k({1, 2, 4, 8});
|
||
const Tensor<float> v({1, 3, 4, 8});
|
||
const Tensor<float> o({1, 4, 4, 8});
|
||
const Tensor<float> dO({1, 4, 4, 8});
|
||
Tensor<float> dQ({1, 4, 4, 8});
|
||
Tensor<float> dK({1, 2, 4, 8});
|
||
Tensor<float> dV({1, 3, 4, 8});
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV), std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardLseWrongRank)
|
||
{
|
||
Tensor<float> q({1, 1, 4, 8});
|
||
Tensor<float> k({1, 1, 4, 8});
|
||
Tensor<float> v({1, 1, 4, 8});
|
||
Tensor<float> o({1, 1, 4, 8});
|
||
Tensor<float> dO({1, 1, 4, 8});
|
||
Tensor<float> dQ({1, 1, 4, 8});
|
||
Tensor<float> dK({1, 1, 4, 8});
|
||
Tensor<float> dV({1, 1, 4, 8});
|
||
Tensor<float> lseWrong({1, 1, 4}); // rank-3 instead of rank-4
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 700);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 701);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 702);
|
||
dO.fillWithValue(1.0f);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV, std::nullopt, &lseWrong),
|
||
std::invalid_argument);
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwd, BackwardLseWrongShape)
|
||
{
|
||
Tensor<float> q({2, 4, 16, 8});
|
||
Tensor<float> k({2, 4, 16, 8});
|
||
Tensor<float> v({2, 4, 16, 8});
|
||
Tensor<float> o({2, 4, 16, 8});
|
||
Tensor<float> dO({2, 4, 16, 8});
|
||
Tensor<float> dQ({2, 4, 16, 8});
|
||
Tensor<float> dK({2, 4, 16, 8});
|
||
Tensor<float> dV({2, 4, 16, 8});
|
||
Tensor<float> lseWrong({2, 4, 8, 1}); // Sq=8 instead of Sq=16
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 710);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 711);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 712);
|
||
dO.fillWithValue(1.0f);
|
||
|
||
CpuFpReferenceSdpa::forward(q, k, v, o);
|
||
|
||
EXPECT_THROW(CpuFpReferenceSdpa::backward(q, k, v, o, dO, dQ, dK, dV, std::nullopt, &lseWrong),
|
||
std::invalid_argument);
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Backward mask tests: full leftBound / rightBound / topLeftAlignment support
|
||
// ---------------------------------------------------------------------------
|
||
//
|
||
// Validates that the new full-signature backward() correctly handles all mask
|
||
// types by comparing analytical gradients against finite-difference numerical
|
||
// gradients computed from the forward pass.
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardCausalMaskExplicitBounds)
|
||
{
|
||
// Verify that backward(..., leftBound=-1, rightBound=0, topLeftAlignment=true)
|
||
// produces identical results to backward(..., causalMask=true).
|
||
// [B=1, H=1, Sq=4, Skv=4, D=4, Dv=4]
|
||
Tensor<float> q({1, 1, 4, 4});
|
||
Tensor<float> k({1, 1, 4, 4});
|
||
Tensor<float> v({1, 1, 4, 4});
|
||
Tensor<float> o({1, 1, 4, 4});
|
||
Tensor<float> dO({1, 1, 4, 4});
|
||
Tensor<float> dQBool({1, 1, 4, 4});
|
||
Tensor<float> dKBool({1, 1, 4, 4});
|
||
Tensor<float> dVBool({1, 1, 4, 4});
|
||
Tensor<float> dQExplicit({1, 1, 4, 4});
|
||
Tensor<float> dKExplicit({1, 1, 4, 4});
|
||
Tensor<float> dVExplicit({1, 1, 4, 4});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 8000);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 8001);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 8002);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 8003);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
CpuFpReferenceSdpa::forward(q, k, v, o, std::nullopt, noMask, /*causalMask=*/true);
|
||
|
||
// Bool convenience overload
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQBool,
|
||
dKBool,
|
||
dVBool,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
/*causalMask=*/true);
|
||
|
||
// Explicit bounds overload (equivalent to TOP_LEFT_CAUSAL)
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQExplicit,
|
||
dKExplicit,
|
||
dVExplicit,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
/*leftBound=*/static_cast<int64_t>(-1),
|
||
/*rightBound=*/static_cast<int64_t>(0),
|
||
/*topLeftAlignment=*/true);
|
||
|
||
const float tol = 1e-6f;
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
EXPECT_NEAR(dQBool.getHostValue(0, 0, sq, d), dQExplicit.getHostValue(0, 0, sq, d), tol)
|
||
<< "dQ mismatch at sq=" << sq << ", d=" << d;
|
||
EXPECT_NEAR(dKBool.getHostValue(0, 0, sq, d), dKExplicit.getHostValue(0, 0, sq, d), tol)
|
||
<< "dK mismatch at sq=" << sq << ", d=" << d;
|
||
EXPECT_NEAR(dVBool.getHostValue(0, 0, sq, d), dVExplicit.getHostValue(0, 0, sq, d), tol)
|
||
<< "dV mismatch at sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, BottomRightCausalMask)
|
||
{
|
||
// BOTTOM_RIGHT_CAUSAL: leftBound=-1, rightBound=0, topLeftAlignment=false
|
||
// [B=1, H=1, Sq=3, Skv=5, D=4, Dv=4] (Skv > Sq to exercise BR offset)
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 3;
|
||
constexpr int64_t SEQ_KV = 5;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 8100);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 8101);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 8102);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 8103);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = -1;
|
||
constexpr int64_t RIGHT_BOUND = 0;
|
||
constexpr bool TOP_LEFT = false;
|
||
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQ,
|
||
dK,
|
||
dV,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (BR causal)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (BR causal)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (BR causal)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, WindowMaskSymmetricTopLeft)
|
||
{
|
||
// WINDOW_GENERIC: symmetric window (leftBound=1, rightBound=1) with top-left alignment
|
||
// [B=1, H=1, Sq=4, Skv=4, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 4;
|
||
constexpr int64_t SEQ_KV = 4;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 8200);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 8201);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 8202);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 8203);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = 1;
|
||
constexpr int64_t RIGHT_BOUND = 1;
|
||
constexpr bool TOP_LEFT = true;
|
||
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQ,
|
||
dK,
|
||
dV,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (window TL sym)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (window TL sym)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (window TL sym)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, WindowMaskAsymmetricBottomRight)
|
||
{
|
||
// WINDOW_GENERIC: asymmetric window (leftBound=2, rightBound=1) with bottom-right alignment
|
||
// [B=1, H=1, Sq=4, Skv=5, D=4, Dv=4] (Skv > Sq to exercise BR offset)
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 4;
|
||
constexpr int64_t SEQ_KV = 5;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 8300);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 8301);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 8302);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 8303);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = 2;
|
||
constexpr int64_t RIGHT_BOUND = 1;
|
||
constexpr bool TOP_LEFT = false;
|
||
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQ,
|
||
dK,
|
||
dV,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (window BR asym)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (window BR asym)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (window BR asym)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, WindowMaskLargerSq)
|
||
{
|
||
// WINDOW_GENERIC with Sq > Skv to test edge where some rows are fully masked
|
||
// [B=1, H=1, Sq=5, Skv=3, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS = 1;
|
||
constexpr int64_t SEQ_Q = 5;
|
||
constexpr int64_t SEQ_KV = 3;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 8400);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 8401);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 8402);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 8403);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = 1;
|
||
constexpr int64_t RIGHT_BOUND = 1;
|
||
constexpr bool TOP_LEFT = false;
|
||
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQ,
|
||
dK,
|
||
dV,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (window BR Sq>Skv)");
|
||
checkGradient(k, dK, {BATCH, HEADS, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (window BR Sq>Skv)");
|
||
checkGradient(v, dV, {BATCH, HEADS, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (window BR Sq>Skv)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, GQABottomRightCausalMask)
|
||
{
|
||
// Combined: GQA + bottom-right causal mask
|
||
// [B=1, H_q=4, H_kv=2, Sq=3, Skv=5, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS_Q = 4;
|
||
constexpr int64_t HEADS_KV = 2;
|
||
constexpr int64_t SEQ_Q = 3;
|
||
constexpr int64_t SEQ_KV = 5;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 8500);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 8501);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 8502);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 8503);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = -1;
|
||
constexpr int64_t RIGHT_BOUND = 0;
|
||
constexpr bool TOP_LEFT = false;
|
||
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQ,
|
||
dK,
|
||
dV,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS_Q, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (GQA+BR causal)");
|
||
checkGradient(k, dK, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (GQA+BR causal)");
|
||
checkGradient(v, dV, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (GQA+BR causal)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaGradCheck, GQAWindowMask)
|
||
{
|
||
// Combined: GQA + window mask (leftBound=1, rightBound=2, topLeftAlignment=true)
|
||
// [B=1, H_q=4, H_kv=1, Sq=4, Skv=4, D=4, Dv=4]
|
||
constexpr int64_t BATCH = 1;
|
||
constexpr int64_t HEADS_Q = 4;
|
||
constexpr int64_t HEADS_KV = 1;
|
||
constexpr int64_t SEQ_Q = 4;
|
||
constexpr int64_t SEQ_KV = 4;
|
||
constexpr int64_t HEAD_DIM = 4;
|
||
constexpr int64_t HEAD_DIM_V = 4;
|
||
|
||
Tensor<float> q({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> k({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> v({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
Tensor<float> o({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
Tensor<float> dO({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM_V});
|
||
|
||
q.fillWithRandomValues(-0.5f, 0.5f, 8600);
|
||
k.fillWithRandomValues(-0.5f, 0.5f, 8601);
|
||
v.fillWithRandomValues(-0.5f, 0.5f, 8602);
|
||
dO.fillWithRandomValues(-0.5f, 0.5f, 8603);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = 1;
|
||
constexpr int64_t RIGHT_BOUND = 2;
|
||
constexpr bool TOP_LEFT = true;
|
||
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
|
||
Tensor<float> dQ({BATCH, HEADS_Q, SEQ_Q, HEAD_DIM});
|
||
Tensor<float> dK({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM});
|
||
Tensor<float> dV({BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V});
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQ,
|
||
dK,
|
||
dV,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
auto fwdLoss = [&]() -> double {
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT);
|
||
return computeDotProductLoss(dO, o);
|
||
};
|
||
|
||
checkGradient(q, dQ, {BATCH, HEADS_Q, SEQ_Q, HEAD_DIM}, fwdLoss, "dQ (GQA+window)");
|
||
checkGradient(k, dK, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM}, fwdLoss, "dK (GQA+window)");
|
||
checkGradient(v, dV, {BATCH, HEADS_KV, SEQ_KV, HEAD_DIM_V}, fwdLoss, "dV (GQA+window)");
|
||
}
|
||
|
||
TEST(TestCpuFpReferenceSdpaBwdFp32, BackwardWithLSEAndWindowMask)
|
||
{
|
||
// Verifies LSE-based backward and scratch backward produce identical results
|
||
// when a window mask is applied.
|
||
// [B=1, H=2, Sq=4, Skv=4, D=4, Dv=4]
|
||
Tensor<float> q({1, 2, 4, 4});
|
||
Tensor<float> k({1, 2, 4, 4});
|
||
Tensor<float> v({1, 2, 4, 4});
|
||
Tensor<float> o({1, 2, 4, 4});
|
||
Tensor<float> lse({1, 2, 4, 1});
|
||
Tensor<float> dO({1, 2, 4, 4});
|
||
Tensor<float> dQWithLse({1, 2, 4, 4});
|
||
Tensor<float> dKWithLse({1, 2, 4, 4});
|
||
Tensor<float> dVWithLse({1, 2, 4, 4});
|
||
Tensor<float> dQNoLse({1, 2, 4, 4});
|
||
Tensor<float> dKNoLse({1, 2, 4, 4});
|
||
Tensor<float> dVNoLse({1, 2, 4, 4});
|
||
|
||
q.fillWithRandomValues(-1.0f, 1.0f, 8700);
|
||
k.fillWithRandomValues(-1.0f, 1.0f, 8701);
|
||
v.fillWithRandomValues(-1.0f, 1.0f, 8702);
|
||
dO.fillWithRandomValues(-1.0f, 1.0f, 8703);
|
||
|
||
const TensorBase<float>* noMask = nullptr;
|
||
constexpr int64_t LEFT_BOUND = 1;
|
||
constexpr int64_t RIGHT_BOUND = 1;
|
||
constexpr bool TOP_LEFT = true;
|
||
|
||
// Forward pass with LSE and window mask
|
||
CpuFpReferenceSdpa::forward(
|
||
q, k, v, o, std::nullopt, noMask, LEFT_BOUND, RIGHT_BOUND, TOP_LEFT, &lse);
|
||
|
||
// Backward WITH LSE
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQWithLse,
|
||
dKWithLse,
|
||
dVWithLse,
|
||
std::nullopt,
|
||
&lse,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
// Backward WITHOUT LSE (recompute from scratch)
|
||
CpuFpReferenceSdpa::backward(q,
|
||
k,
|
||
v,
|
||
o,
|
||
dO,
|
||
dQNoLse,
|
||
dKNoLse,
|
||
dVNoLse,
|
||
std::nullopt,
|
||
nullptr,
|
||
noMask,
|
||
LEFT_BOUND,
|
||
RIGHT_BOUND,
|
||
TOP_LEFT);
|
||
|
||
const float tol = 1e-5f;
|
||
for(int h = 0; h < 2; ++h)
|
||
{
|
||
for(int sq = 0; sq < 4; ++sq)
|
||
{
|
||
for(int d = 0; d < 4; ++d)
|
||
{
|
||
EXPECT_NEAR(
|
||
dQWithLse.getHostValue(0, h, sq, d), dQNoLse.getHostValue(0, h, sq, d), tol)
|
||
<< "dQ mismatch at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
EXPECT_NEAR(
|
||
dKWithLse.getHostValue(0, h, sq, d), dKNoLse.getHostValue(0, h, sq, d), tol)
|
||
<< "dK mismatch at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
EXPECT_NEAR(
|
||
dVWithLse.getHostValue(0, h, sq, d), dVNoLse.getHostValue(0, h, sq, d), tol)
|
||
<< "dV mismatch at h=" << h << ", sq=" << sq << ", d=" << d;
|
||
}
|
||
}
|
||
}
|
||
}
|