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();