From fef595d077e9c1382e6aee555679f0a3ea3f4d51 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Tue, 3 Nov 2026 05:16:36 -0800 Subject: [PATCH] Accumulate generic tile flash attention directly into att_out. att_out is now an IN/OUT parameter: the generic TileFlashAttentionReturnExpSumsAndMaxLogits kernel uses it directly as the online-softmax accumulator instead of a separate C accumulator buffer. - Remove the C accumulator region from TileFlashAttentionWorkspaceLayout (saves padded_q_count * qkv_dim * 4 bytes of per-worker workspace). - Drop the per-call zero fill of C and the final C -> att_out copy. - Address att_out via Row()/Stride() so padded layouts work. - Document att_out as in-out, like exp_denominator_sums and max_logits: it must be zero-initialized before the first call, and its rows must be aligned to the SIMD vector size (all current callers already satisfy both). Fix the MixedRanges test and the attention benchmark, which passed uninitialized att_out. - Add TestTiledFlashAttentionChainedCalls with a kernel-agnostic helper, CheckChainedCallsMatchSingleCall, verifying that chaining calls over sub-ranges of the KV (carrying att_out, exp_sums and max_logits) matches a single call over the full range, including remainder query groups. The MixedRanges and chaining tests share a TiledFlashAttentionTestSetup. PiperOrigin-RevId: 1012927185 --- gemma/flash_attention.cc | 47 ++--- gemma/flash_attention_test.cc | 320 ++++++++++++++++++++++++++-------- 2 files changed, 262 insertions(+), 105 deletions(-) diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index 112929af..94b690c6 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -1099,7 +1099,12 @@ 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 of shape (q_count, qkv_dim) +// att_out - MatPtrT of shape (q_count, qkv_dim), used directly as the +// online-softmax accumulator via aligned loads/stores, so each row must be +// aligned to the SIMD vector size. It is in-out: initialize it to zero before +// the first call. Rows of queries that attend to no position are left +// unchanged, so it can be carried across calls over successive kv chunks +// together with exp_denominator_sums and max_logits. // exp_denominator_sums and max_logits - float* of shape (q_count,), holding // the running online-softmax denominator and maximum logit per query. Only // entries [0, q_count) are read and written. They are in-out: initialize them @@ -1485,8 +1490,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& att_out, + size_t kBlockSize, FlashStepBufT* HWY_RESTRICT step_consts_buf, const KV_T** HWY_RESTRICT step_v_tiles, const BF16** HWY_RESTRICT step_microscaling_v_ptrs, @@ -1552,8 +1557,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* dst = step_consts_buf + @@ -1585,7 +1590,9 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle( } MatPtrT 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()) { MulByConstAndAddTileUpTo8( @@ -1609,7 +1616,6 @@ struct TileFlashAttentionWorkspaceLayout { using StepBufT = FlashStepBufT; 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; @@ -1620,15 +1626,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), @@ -1696,7 +1698,7 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( } TileFlashAttentionWorkspaceLayout 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. @@ -1713,8 +1715,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 = @@ -1732,9 +1732,6 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( float* step_q_scales_s_buf = HWY_RCAST_ALIGNED(float*, raw_ptr + layout.step_qb_offset); - hwy::Span 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 softmax_buf_span(softmax_buf_ptr, padded_q_count * kBlockSize); hwy::Span min_start_pos_per_group(pos_data_ptr, num_loops); @@ -1784,7 +1781,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); }; @@ -1827,16 +1824,6 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( current_kv_idx++; } } - - using VF = hn::Vec; - 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 diff --git a/gemma/flash_attention_test.cc b/gemma/flash_attention_test.cc index 8afec943..85589186 100644 --- a/gemma/flash_attention_test.cc +++ b/gemma/flash_attention_test.cc @@ -861,78 +861,76 @@ void TestTiledFlashAttentionInt8MatrixAccumulation() { enum class MixedRangesQueryType { kF32, kBF16, kInt16, kInt8 }; -// Queries in the same kernel group (4 or 8 queries) whose [start, last] ranges -// cover different kernel blocks: some start late, some end early. A query that -// is inactive in a block must not pick up stale step constants from the -// workspace, which is poisoned with 0xFF (NaN) here. The reference runs every -// query on its own, so no query is ever inactive while others are active. +// Shared setup for the generic (non-AMX, non-matrix accumulation) tile flash +// attention kernels: a Gaussian KV cache of `kSeqLen` tokens (more than two +// 256-token kernel blocks, with a partial last tile), `num_queries` Gaussian +// queries compressed to `q_type`, and per-query [start, last] ranges mixing +// full, late-start, early-end and short ranges within each kernel group. template -void RunTiledFlashAttentionMixedRangesTest(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("kv", Extents2D(num_tiles, tile_elems), ctx.allocator, - MatPadding::kPacked); - PopulateTestKVCache(kv, kv_encoding, kQkvDim, - TestDataDistribution::kGaussian); - const hwy::Span kvs(&kv, 1); - - const AlignedFloatVector q_f32 = PopulateTestQueries( - kNumQueries, kQkvDim, TestDataDistribution::kGaussian); - std::vector> q_bf16(kNumQueries * kQkvDim); - std::vector> q_i16(kNumQueries * - kQkvDim); - std::vector> 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()); - } - - std::vector> start_pos(kNumQueries); - std::vector> last_pos(kNumQueries); - for (size_t i = 0; i < kNumQueries; ++i) { - switch (i % 4) { - case 0: // Full range. - start_pos[i] = 0; - last_pos[i] = kSeqLen - 1; - break; - case 1: // Late start: inactive in the first block. - start_pos[i] = 300; - last_pos[i] = kSeqLen - 1; - break; - case 2: // Early end: inactive in the later blocks. - start_pos[i] = 0; - last_pos[i] = 200; - break; - default: // Short range spanning a block boundary. - start_pos[i] = 250; - last_pos[i] = 270; - break; +struct TiledFlashAttentionTestSetup { + static constexpr size_t kQkvDim = 64; + static constexpr size_t kSeqLen = 600; + using PosVec = std::vector>; + + TiledFlashAttentionTestSetup(gcpp::KVEncoding kv_encoding, + MixedRangesQueryType q_type, size_t num_queries, + ThreadingContext& ctx) + : q_type(q_type), + num_queries(num_queries), + kv("kv", + Extents2D(hwy::DivCeil(kSeqLen, gcpp::KVCache::kTileSize), + *gcpp::GetTileSizeBytes(kv_encoding, kQkvDim) / + sizeof(KV_T)), + ctx.allocator, MatPadding::kPacked), + q_f32(PopulateTestQueries(num_queries, kQkvDim, + TestDataDistribution::kGaussian)), + q_bf16(num_queries * kQkvDim), + q_i16(num_queries * kQkvDim), + q_i8(num_queries * kQkvDim), + q_scales(num_queries), + start_pos(num_queries), + last_pos(num_queries) { + PopulateTestKVCache(kv, kv_encoding, kQkvDim, + TestDataDistribution::kGaussian); + if (q_type == MixedRangesQueryType::kBF16) { + CompressQueriesBF16Contiguous(q_f32.data(), kQkvDim, num_queries, + q_bf16.data()); + } else if (q_type == MixedRangesQueryType::kInt16) { + CompressQueriesInt16Contiguous(q_f32.data(), kQkvDim, num_queries, + q_i16.data(), q_scales.data()); + } else if (q_type == MixedRangesQueryType::kInt8) { + CompressQueriesUint8Contiguous(q_f32.data(), kQkvDim, num_queries, + q_i8.data(), q_scales.data()); + } + for (size_t i = 0; i < num_queries; ++i) { + switch (i % 4) { + case 0: // Full range. + start_pos[i] = 0; + last_pos[i] = kSeqLen - 1; + break; + case 1: // Late start: inactive in the first block. + start_pos[i] = 300; + last_pos[i] = kSeqLen - 1; + break; + case 2: // Early end: inactive in the later blocks. + start_pos[i] = 0; + last_pos[i] = 200; + break; + default: // Short range spanning a block boundary. + start_pos[i] = 250; + last_pos[i] = 270; + break; + } } } - // Runs `q_count` queries starting at `first_q`. - auto run = [&](size_t first_q, size_t q_count, - hwy::Span starts, hwy::Span lasts, - MatPtrT& att_out, float* exp_sums, float* max_logits, - hwy::AlignedVector* workspace) { + // Runs the kernel for `q_count` queries starting at `first_q`, with the + // given per-query ranges, updating att_out, exp_sums and max_logits in place. + void Run(size_t first_q, size_t q_count, hwy::Span starts, + hwy::Span lasts, MatPtrT& att_out, + float* exp_sums, float* max_logits, + hwy::AlignedVector* workspace) const { + const hwy::Span kvs(&kv, 1); const hwy::Span scales(q_scales.data() + first_q, q_count); switch (q_type) { case MixedRangesQueryType::kF32: @@ -956,7 +954,37 @@ void RunTiledFlashAttentionMixedRangesTest(gcpp::KVEncoding kv_encoding, lasts, /*att_cap=*/0.0f, att_out, exp_sums, max_logits, workspace); break; } - }; + } + + const MixedRangesQueryType q_type; + const size_t num_queries; + MatStorageT kv; + AlignedFloatVector q_f32; + std::vector> q_bf16; + std::vector> q_i16; + std::vector> q_i8; + AlignedFloatVector q_scales; + PosVec start_pos; + PosVec last_pos; +}; + +// Queries in the same kernel group (4 or 8 queries) whose [start, last] ranges +// cover different kernel blocks: some start late, some end early. A query that +// is inactive in a block must not pick up stale step constants from the +// workspace, which is poisoned with 0xFF (NaN) here. The reference runs every +// query on its own, so no query is ever inactive while others are active. +template +void RunTiledFlashAttentionMixedRangesTest(gcpp::KVEncoding kv_encoding, + MixedRangesQueryType q_type, + const char* name, float tol) { + using Setup = TiledFlashAttentionTestSetup; + constexpr size_t kQkvDim = Setup::kQkvDim; + constexpr size_t kNumQueries = 16; + ThreadingArgs threading_args; + ThreadingContext ctx(threading_args); + const Setup setup(kv_encoding, q_type, kNumQueries, ctx); + const auto& start_pos = setup.start_pos; + const auto& last_pos = setup.last_pos; constexpr float kInitMax = -std::numeric_limits::max() / 2.0f; const size_t padded_queries = hwy::RoundUpTo(kNumQueries, 16); @@ -964,23 +992,25 @@ void RunTiledFlashAttentionMixedRangesTest(gcpp::KVEncoding kv_encoding, // All queries in one call, with a NaN-poisoned workspace. MatStorageT 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 exp_sums(padded_queries, 0.0f); std::vector max_logits(padded_queries, kInitMax); hwy::AlignedVector poisoned_workspace(size_t{1} << 20, 0xFF); - run(0, kNumQueries, hwy::Span(start_pos), - hwy::Span(last_pos), att_out, exp_sums.data(), - max_logits.data(), &poisoned_workspace); + setup.Run(0, kNumQueries, hwy::Span(start_pos), + hwy::Span(last_pos), att_out, exp_sums.data(), + max_logits.data(), &poisoned_workspace); size_t num_failed_queries = 0; for (size_t i = 0; i < kNumQueries; ++i) { MatStorageT ref_out("ref_out", Extents2D(1, kQkvDim), ctx.allocator, MatPadding::kPacked); + ZeroInit(ref_out); std::vector ref_exp_sums(16, 0.0f); std::vector ref_max_logits(16, kInitMax); hwy::AlignedVector ref_workspace(size_t{1} << 20, 0); - run(i, 1, hwy::Span(&start_pos[i], 1), - hwy::Span(&last_pos[i], 1), ref_out, ref_exp_sums.data(), - ref_max_logits.data(), &ref_workspace); + setup.Run(i, 1, hwy::Span(&start_pos[i], 1), + hwy::Span(&last_pos[i], 1), ref_out, + ref_exp_sums.data(), ref_max_logits.data(), &ref_workspace); // The single-query reference starts its block loop at its own start // position, so block boundaries (and hence rescaling and BF16/quantized @@ -1029,6 +1059,145 @@ 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 +void CheckChainedCallsMatchSingleCall(const KernelFn& kernel, + ThreadingContext& ctx, + size_t num_queries, size_t qkv_dim, + hwy::Span starts, + hwy::Span lasts, + size_t num_chunks, const char* name, + float tol) { + using PosVec = std::vector>; + constexpr float kInitMax = -std::numeric_limits::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 single_out("single_out", Extents2D(num_queries, qkv_dim), + ctx.allocator, MatPadding::kPacked); + ZeroInit(single_out); + std::vector single_sums(padded_queries, 0.0f); + std::vector single_max(padded_queries, kInitMax); + hwy::AlignedVector workspace(kWorkspaceBytes, 0); + kernel(starts, lasts, single_out, single_sums.data(), single_max.data(), + &workspace); + + // Chained calls over consecutive sub-ranges. + MatStorageT chained_out("chained_out", + Extents2D(num_queries, qkv_dim), + ctx.allocator, MatPadding::kOdd); + ZeroInit(chained_out); + std::vector chained_sums(padded_queries, 0.0f); + std::vector 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; + // Every chunk must be non-empty for every query: the [start, last] API + // cannot express an empty range. + 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 poisoned_workspace(kWorkspaceBytes, 0xFF); + kernel(hwy::Span(chunk_starts), + hwy::Span(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 tile flash attention +// kernels. 15 queries also exercise the remainder groups (8+4+2+1 or +// 4+4+4+2+1), which now write directly into the unpadded att_out rows. +template +void RunTiledFlashAttentionChainingTest(gcpp::KVEncoding kv_encoding, + MixedRangesQueryType q_type, + const char* name, float tol) { + using Setup = TiledFlashAttentionTestSetup; + ThreadingArgs threading_args; + ThreadingContext ctx(threading_args); + for (size_t num_queries : {size_t{16}, size_t{15}}) { + const Setup setup(kv_encoding, q_type, num_queries, ctx); + auto kernel = [&](hwy::Span starts, + hwy::Span lasts, MatPtrT& att_out, + float* exp_sums, float* max_logits, + hwy::AlignedVector* workspace) { + setup.Run(0, num_queries, starts, lasts, att_out, exp_sums, max_logits, + workspace); + }; + for (size_t num_chunks : {size_t{2}, size_t{3}}) { + CheckChainedCallsMatchSingleCall( + kernel, ctx, num_queries, Setup::kQkvDim, + hwy::Span(setup.start_pos), + hwy::Span(setup.last_pos), num_chunks, name, tol); + } + } +} + +void TestTiledFlashAttentionChainedCalls() { + RunTiledFlashAttentionChainingTest( + 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( + gcpp::KVEncoding::kBF16TwoTranspositions, MixedRangesQueryType::kBF16, + "BF16TTkv_BF16q", 5e-3f); + RunTiledFlashAttentionChainingTest( + gcpp::KVEncoding::kInt8, MixedRangesQueryType::kF32, "Int8kv_F32q", + 5e-3f); + RunTiledFlashAttentionChainingTest( + gcpp::KVEncoding::kInt8TwoTranspositions, MixedRangesQueryType::kBF16, + "Int8TTkv_BF16q", 5e-3f); + RunTiledFlashAttentionChainingTest( + gcpp::KVEncoding::kInt8TwoTranspositions, MixedRangesQueryType::kInt16, + "Int8TTkv_Int16q", 5e-3f); + RunTiledFlashAttentionChainingTest( + gcpp::KVEncoding::kInt8VNNITwoTranspositions, MixedRangesQueryType::kInt8, + "Int8VNNIkv_Uint8q", 5e-3f); +} + template void RunTiledFlashAttentionDifferentialTest( size_t kv_seq_len, float tol, float tol_exp, float tol_max, @@ -1373,6 +1542,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();