Skip to content

Accumulate generic tile flash attention directly into att_out. #1065

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Accumulate generic tile flash attention directly into att_out. #1065

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 13 additions & 30 deletions gemma/flash_attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1099,7 +1099,8 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidthBF16(
// queries_per_timestep - how many queries begin/end on the same timestep
// attention_shape - see struct definition for more details.
// att_cap - soft cap on attention logits
// att_out - MatPtrT<float> of shape (q_count, qkv_dim)
// att_out - IN/OUT MatPtrT<float> of shape (q_count, qkv_dim). Used directly
// as the accumulator, so it must be zero-initialized before the first call.
// exp_denominator_sums and max_logits: float* of shape:
// (RountedUpTo(q_count,4),)
// Need to be have multiple of 4 elements alocated and
Expand Down Expand Up @@ -1484,8 +1485,8 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle(
const size_t* HWY_RESTRICT min_start_pos_per_group,
const size_t* HWY_RESTRICT max_last_pos_per_group,
const float* HWY_RESTRICT softmax_buf_ptr, float* HWY_RESTRICT max_logits,
float* HWY_RESTRICT exp_denominator_sums,
float* HWY_RESTRICT C_accumulators_ptr, size_t kBlockSize,
float* HWY_RESTRICT exp_denominator_sums, MatPtrT<float>& att_out,
size_t kBlockSize,
FlashStepBufT<Q_T>* HWY_RESTRICT step_consts_buf,
const KV_T** HWY_RESTRICT step_v_tiles,
const BF16** HWY_RESTRICT step_microscaling_v_ptrs,
Expand Down Expand Up @@ -1551,8 +1552,8 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle(
// This query attends to no token in this block, but the multiply-add
// below still runs over all kNumQueries rows with scales_old == 1.
// Zero its step constants (and Int16/Int8 query scales) so that its
// accumulator row is left unchanged; otherwise stale or uninitialized
// workspace values leak into att_out.
// att_out row is left unchanged; otherwise stale step constants from a
// previous group leak into att_out.
constexpr size_t kMaxLanes = hn::MaxLanes(df);
for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) {
FlashStepBufT<Q_T>* dst = step_consts_buf +
Expand Down Expand Up @@ -1584,7 +1585,9 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle(
}

MatPtrT<float> group_out("group_out", Extents2D(kNumQueriesPerLoop, qkv_dim));
group_out.SetPtr(C_accumulators_ptr + q_base_idx * qkv_dim, qkv_dim);
// Accumulate directly into att_out (IN/OUT): rows are rescaled by
// scales_old and the new P*V contribution is added in place.
group_out.SetPtr(att_out.Row(q_base_idx), att_out.Stride());

if constexpr (IsF32<Q_T>()) {
MulByConstAndAddTileUpTo8<kNumQueries>(
Expand All @@ -1608,7 +1611,6 @@ struct TileFlashAttentionWorkspaceLayout {
using StepBufT = FlashStepBufT<Q_T>;

size_t padded_q_count;
size_t c_accum_offset, c_accum_bytes;
size_t softmax_offset, softmax_bytes;
size_t pos_data_offset, pos_data_bytes;
size_t step_consts_offset, step_consts_bytes;
Expand All @@ -1619,15 +1621,11 @@ struct TileFlashAttentionWorkspaceLayout {
size_t step_qb_offset, step_qb_bytes;
size_t total_bytes;

constexpr TileFlashAttentionWorkspaceLayout(size_t q_count, size_t qkv_dim,
size_t kBlockSize,
constexpr TileFlashAttentionWorkspaceLayout(size_t q_count, size_t kBlockSize,
size_t num_loops)
: padded_q_count(hwy::RoundUpTo(
q_count, HWY_MAX(size_t{16}, size_t(kNumQueriesPerLoop)))),
c_accum_offset(0),
c_accum_bytes(hwy::RoundUpTo(padded_q_count * qkv_dim * sizeof(float),
HWY_ALIGNMENT)),
softmax_offset(c_accum_offset + c_accum_bytes),
softmax_offset(0),
softmax_bytes(hwy::RoundUpTo(
padded_q_count * kBlockSize * sizeof(float), HWY_ALIGNMENT)),
pos_data_offset(softmax_offset + softmax_bytes),
Expand Down Expand Up @@ -1695,7 +1693,7 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl(
}

TileFlashAttentionWorkspaceLayout<kNumQueriesPerLoop, KV_T, Q_T> layout(
q_count, qkv_dim, kBlockSize, num_loops);
q_count, kBlockSize, num_loops);
const size_t padded_q_count = layout.padded_q_count;

// Use pre-allocated worker_workspace when available, resizing up if needed.
Expand All @@ -1712,8 +1710,6 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl(
raw_ptr = workspace_fallback.get();
}

float* C_accumulators_ptr =
HWY_RCAST_ALIGNED(float*, raw_ptr + layout.c_accum_offset);
float* softmax_buf_ptr =
HWY_RCAST_ALIGNED(float*, raw_ptr + layout.softmax_offset);
size_t* pos_data_ptr =
Expand All @@ -1731,9 +1727,6 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl(
float* step_q_scales_s_buf =
HWY_RCAST_ALIGNED(float*, raw_ptr + layout.step_qb_offset);

hwy::Span<float> c_accumulators_span(C_accumulators_ptr,
padded_q_count * qkv_dim);
std::fill(c_accumulators_span.begin(), c_accumulators_span.end(), 0.0f);
hwy::Span<float> softmax_buf_span(softmax_buf_ptr,
padded_q_count * kBlockSize);
hwy::Span<size_t> min_start_pos_per_group(pos_data_ptr, num_loops);
Expand Down Expand Up @@ -1783,7 +1776,7 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl(
step_size, att_cap, one_over_cap, start_pos_per_query.data(),
last_pos_per_query.data(), min_start_pos_per_group.data(),
max_last_pos_per_group.data(), softmax_buf_ptr, max_logits,
exp_denominator_sums, C_accumulators_ptr, kBlockSize, step_consts_buf,
exp_denominator_sums, att_out, kBlockSize, step_consts_buf,
step_v_tiles, step_microscaling_v_ptrs, step_microscaling_k_ptrs,
step_q_scales_s_ptrs, step_q_scales_s_buf);
};
Expand Down Expand Up @@ -1826,16 +1819,6 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl(
current_kv_idx++;
}
}

using VF = hn::Vec<DF>;
for (size_t qi = 0; qi < q_count; ++qi) {
float* out = att_out.Row(qi);
const float* accum = C_accumulators_ptr + qi * qkv_dim;
for (size_t d = 0; d < qkv_dim; d += hn::Lanes(df)) {
VF v = hn::LoadU(df, accum + d);
hn::StoreU(v, df, out + d);
}
}
}

template <typename KV_T, typename Q_T>
Expand Down
214 changes: 214 additions & 0 deletions gemma/flash_attention_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include <memory>
#include <numeric>
#include <random>
#include <string>
#include <vector>

#include "compression/types.h"
Expand Down Expand Up @@ -964,6 +965,7 @@ void RunTiledFlashAttentionMixedRangesTest(gcpp::KVEncoding kv_encoding,
// All queries in one call, with a NaN-poisoned workspace.
MatStorageT<float> att_out("att_out", Extents2D(kNumQueries, kQkvDim),
ctx.allocator, MatPadding::kPacked);
ZeroInit(att_out); // att_out is IN/OUT and must start zeroed.
std::vector<float> exp_sums(padded_queries, 0.0f);
std::vector<float> max_logits(padded_queries, kInitMax);
hwy::AlignedVector<uint8_t> poisoned_workspace(size_t{1} << 20, 0xFF);
Expand All @@ -975,6 +977,7 @@ void RunTiledFlashAttentionMixedRangesTest(gcpp::KVEncoding kv_encoding,
for (size_t i = 0; i < kNumQueries; ++i) {
MatStorageT<float> ref_out("ref_out", Extents2D(1, kQkvDim), ctx.allocator,
MatPadding::kPacked);
ZeroInit(ref_out);
std::vector<float> ref_exp_sums(16, 0.0f);
std::vector<float> ref_max_logits(16, kInitMax);
hwy::AlignedVector<uint8_t> ref_workspace(size_t{1} << 20, 0);
Expand Down Expand Up @@ -1029,6 +1032,216 @@ void TestTiledFlashAttentionMixedQueryRanges() {
"Int8VNNIkv_Uint8q", 5e-3f);
}

// Kernel-agnostic check that tile flash attention kernels can be chained over
// the KV range: running `kernel` once over each query's full
// [starts[i], lasts[i]] range must match `num_chunks` consecutive calls over
// sub-ranges that carry att_out, exp_sums and max_logits between calls.
//
// `kernel(starts, lasts, att_out, exp_sums, max_logits, workspace)` must run
// all `num_queries` queries and update att_out (IN/OUT), exp_sums and
// max_logits in place. Each query is split at its own fractional points, so
// chunk boundaries differ across queries and do not align with kernel blocks.
// The chained calls use a fresh 0xFF-poisoned (NaN) workspace per call, so no
// state may be kept in the workspace, and a padded (kOdd) att_out, so att_out
// must be addressed via its stride.
template <class KernelFn>
void CheckChainedCallsMatchSingleCall(const KernelFn& kernel,
ThreadingContext& ctx,
size_t num_queries, size_t qkv_dim,
hwy::Span<const size_t> starts,
hwy::Span<const size_t> lasts,
size_t num_chunks, const std::string& name,
float tol) {
using PosVec = std::vector<size_t, hwy::AlignedAllocator<size_t>>;
constexpr float kInitMax = -std::numeric_limits<float>::max() / 2.0f;
constexpr size_t kWorkspaceBytes = size_t{1} << 20;
const size_t padded_queries = hwy::RoundUpTo(num_queries, 16);

// Single call over the full ranges.
MatStorageT<float> single_out("single_out", Extents2D(num_queries, qkv_dim),
ctx.allocator, MatPadding::kPacked);
ZeroInit(single_out);
std::vector<float> single_sums(padded_queries, 0.0f);
std::vector<float> single_max(padded_queries, kInitMax);
hwy::AlignedVector<uint8_t> workspace(kWorkspaceBytes, 0);
kernel(starts, lasts, single_out, single_sums.data(), single_max.data(),
&workspace);

// Chained calls over consecutive sub-ranges.
MatStorageT<float> chained_out("chained_out",
Extents2D(num_queries, qkv_dim),
ctx.allocator, MatPadding::kOdd);
ZeroInit(chained_out);
std::vector<float> chained_sums(padded_queries, 0.0f);
std::vector<float> chained_max(padded_queries, kInitMax);
PosVec chunk_starts(num_queries);
PosVec chunk_lasts(num_queries);
for (size_t c = 0; c < num_chunks; ++c) {
for (size_t i = 0; i < num_queries; ++i) {
const size_t len = lasts[i] - starts[i] + 1;
HWY_ASSERT(len >= num_chunks);
chunk_starts[i] = starts[i] + len * c / num_chunks;
chunk_lasts[i] = starts[i] + len * (c + 1) / num_chunks - 1;
}
hwy::AlignedVector<uint8_t> poisoned_workspace(kWorkspaceBytes, 0xFF);
kernel(hwy::Span<const size_t>(chunk_starts),
hwy::Span<const size_t>(chunk_lasts), chained_out,
chained_sums.data(), chained_max.data(), &poisoned_workspace);
}

// Block boundaries (and hence rescaling and BF16/quantized rounding of the
// step constants) differ between the single and chained calls, so compare
// with tolerances.
size_t num_failed_queries = 0;
for (size_t i = 0; i < num_queries; ++i) {
float max_out_err = 0.0f;
bool all_finite = true;
for (size_t j = 0; j < qkv_dim; ++j) {
const float out = chained_out.Row(i)[j];
all_finite &= std::isfinite(out);
max_out_err =
std::max(max_out_err, std::abs(out - single_out.Row(i)[j]));
}
const float sum_err = std::abs(chained_sums[i] - single_sums[i]) /
std::max(1.0f, std::abs(single_sums[i]));
const float max_err = std::abs(chained_max[i] - single_max[i]);
if (!all_finite || !(max_out_err <= tol) || !(sum_err <= tol) ||
!(max_err <= tol)) {
++num_failed_queries;
ADD_FAILURE() << name << " chunks=" << num_chunks << " q=" << i << " ["
<< starts[i] << "," << lasts[i] << "] finite=" << all_finite
<< " max_att_out_err=" << max_out_err
<< " rel_exp_sum_err=" << sum_err
<< " max_logit_err=" << max_err << " tol=" << tol;
}
}
EXPECT_EQ(num_failed_queries, size_t{0}) << name << " chunks=" << num_chunks;
}

// Runs CheckChainedCallsMatchSingleCall on the generic (non-AMX, non-matrix
// accumulation) tile flash attention kernels.
template <typename KV_T>
void RunTiledFlashAttentionChainingTest(gcpp::KVEncoding kv_encoding,
MixedRangesQueryType q_type,
const char* name, float tol) {
constexpr size_t kQkvDim = 64;
// More than two kernel blocks (256 tokens), with a partial last tile.
constexpr size_t kSeqLen = 600;
constexpr size_t kNumQueries = 16;
ThreadingArgs threading_args;
ThreadingContext ctx(threading_args);

const size_t num_tiles = hwy::DivCeil(kSeqLen, gcpp::KVCache::kTileSize);
const size_t tile_elems =
*gcpp::GetTileSizeBytes(kv_encoding, kQkvDim) / sizeof(KV_T);
MatStorageT<KV_T> kv("kv", Extents2D(num_tiles, tile_elems), ctx.allocator,
MatPadding::kPacked);
PopulateTestKVCache(kv, kv_encoding, kQkvDim,
TestDataDistribution::kGaussian);
const hwy::Span<const MatPtr> kvs(&kv, 1);

const AlignedFloatVector q_f32 = PopulateTestQueries(
kNumQueries, kQkvDim, TestDataDistribution::kGaussian);
std::vector<BF16, hwy::AlignedAllocator<BF16>> q_bf16(kNumQueries * kQkvDim);
std::vector<int16_t, hwy::AlignedAllocator<int16_t>> q_i16(kNumQueries *
kQkvDim);
std::vector<int8_t, hwy::AlignedAllocator<int8_t>> q_i8(kNumQueries *
kQkvDim);
AlignedFloatVector q_scales(kNumQueries);
if (q_type == MixedRangesQueryType::kBF16) {
CompressQueriesBF16Contiguous(q_f32.data(), kQkvDim, kNumQueries,
q_bf16.data());
} else if (q_type == MixedRangesQueryType::kInt16) {
CompressQueriesInt16Contiguous(q_f32.data(), kQkvDim, kNumQueries,
q_i16.data(), q_scales.data());
} else if (q_type == MixedRangesQueryType::kInt8) {
CompressQueriesUint8Contiguous(q_f32.data(), kQkvDim, kNumQueries,
q_i8.data(), q_scales.data());
}
const hwy::Span<const float> scales(q_scales.data(), kNumQueries);

// Mix of full, late-start, early-end and short ranges within each group.
std::vector<size_t, hwy::AlignedAllocator<size_t>> start_pos(kNumQueries);
std::vector<size_t, hwy::AlignedAllocator<size_t>> last_pos(kNumQueries);
for (size_t i = 0; i < kNumQueries; ++i) {
switch (i % 4) {
case 0:
start_pos[i] = 0;
last_pos[i] = kSeqLen - 1;
break;
case 1:
start_pos[i] = 300;
last_pos[i] = kSeqLen - 1;
break;
case 2:
start_pos[i] = 0;
last_pos[i] = 200;
break;
default:
start_pos[i] = 250;
last_pos[i] = 270;
break;
}
}

auto kernel = [&](hwy::Span<const size_t> starts,
hwy::Span<const size_t> lasts, MatPtrT<float>& att_out,
float* exp_sums, float* max_logits,
hwy::AlignedVector<uint8_t>* workspace) {
switch (q_type) {
case MixedRangesQueryType::kF32:
DispatchTileFlashAttentionReturnExpSumsAndMaxLogits(
kvs, kNumQueries, q_f32.data(), starts, lasts, /*att_cap=*/0.0f,
att_out, exp_sums, max_logits, workspace);
break;
case MixedRangesQueryType::kBF16:
DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsBF16(
kvs, kNumQueries, q_bf16.data(), starts, lasts, /*att_cap=*/0.0f,
att_out, exp_sums, max_logits, workspace);
break;
case MixedRangesQueryType::kInt16:
DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16(
kvs, kNumQueries, q_i16.data(), scales, starts, lasts,
/*att_cap=*/0.0f, att_out, exp_sums, max_logits, workspace);
break;
case MixedRangesQueryType::kInt8:
DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8(
kvs, kNumQueries, q_i8.data(), scales, starts, lasts,
/*att_cap=*/0.0f, att_out, exp_sums, max_logits, workspace);
break;
}
};

for (size_t num_chunks : {size_t{2}, size_t{3}}) {
CheckChainedCallsMatchSingleCall(
kernel, ctx, kNumQueries, kQkvDim, hwy::Span<const size_t>(start_pos),
hwy::Span<const size_t>(last_pos), num_chunks, name, tol);
}
}

void TestTiledFlashAttentionChainedCalls() {
RunTiledFlashAttentionChainingTest<float>(
gcpp::KVEncoding::kF32, MixedRangesQueryType::kF32, "F32kv_F32q", 1e-4f);
// BF16 step constants are normalized by a different partial sum in each
// chained call, so their rounding differs from the single call; short
// ranges do not average this out, hence the looser tolerance.
RunTiledFlashAttentionChainingTest<BF16>(
gcpp::KVEncoding::kBF16TwoTranspositions, MixedRangesQueryType::kBF16,
"BF16TTkv_BF16q", 5e-3f);
RunTiledFlashAttentionChainingTest<int8_t>(
gcpp::KVEncoding::kInt8, MixedRangesQueryType::kF32, "Int8kv_F32q",
5e-3f);
RunTiledFlashAttentionChainingTest<int8_t>(
gcpp::KVEncoding::kInt8TwoTranspositions, MixedRangesQueryType::kBF16,
"Int8TTkv_BF16q", 5e-3f);
RunTiledFlashAttentionChainingTest<int8_t>(
gcpp::KVEncoding::kInt8TwoTranspositions, MixedRangesQueryType::kInt16,
"Int8TTkv_Int16q", 5e-3f);
RunTiledFlashAttentionChainingTest<int8_t>(
gcpp::KVEncoding::kInt8VNNITwoTranspositions, MixedRangesQueryType::kInt8,
"Int8VNNIkv_Uint8q", 5e-3f);
}

template <typename KV_T>
void RunTiledFlashAttentionDifferentialTest(
size_t kv_seq_len, float tol, float tol_exp, float tol_max,
Expand Down Expand Up @@ -1373,6 +1586,7 @@ HWY_EXPORT_AND_TEST_P(FlashAttentionTest,
TestTiledFlashAttentionInt8MatrixAccumulation);
HWY_EXPORT_AND_TEST_P(FlashAttentionTest,
TestTiledFlashAttentionMixedQueryRanges);
HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestTiledFlashAttentionChainedCalls);

HWY_AFTER_TEST();

Expand Down
Loading
Loading