Skip to content

Reduce local-attention KV storage with per-layer rings #1017

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

Closed
Mikyx-1 wants to merge 1 commit into from
Closed

Reduce local-attention KV storage with per-layer rings #1017

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
12 changes: 12 additions & 0 deletions BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,18 @@ cc_library(
],
)

cc_test(
name = "kv_cache_test",
srcs = ["gemma/kv_cache_test.cc"],
deps = [
":configs",
":gemma_args",
":kv_cache",
":threading_context",
"@googletest//:gtest_main",
],
)

cc_test(
name = "flash_attention_test",
srcs = ["gemma/flash_attention_test.cc"],
Expand Down
1 change: 1 addition & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ set(GEMMA_TEST_FILES
compression/sfp_test.cc
evals/gemma_test.cc
gemma/flash_attention_test.cc
gemma/kv_cache_test.cc
gemma/tensor_info_test.cc
io/blob_store_test.cc
io/fields_test.cc
Expand Down
47 changes: 24 additions & 23 deletions gemma/attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,16 @@ namespace gcpp {
namespace HWY_NAMESPACE {

// Computes Q.K scores, which are "logits" (or scores) stored to att.
// `k` is a strided view of the kv cache with dimensions [seq_len, qkv_dim].
// `k` is a strided view of the kv cache with dimensions [cache_rows, qkv_dim].
static HWY_INLINE void QDotK(const size_t start_pos, const size_t last_pos,
const hwy::Divisor& div_seq_len,
const float* HWY_RESTRICT q,
const MatPtrT<KV_t>& k, float* HWY_RESTRICT att,
ThreadingContext& ctx, const size_t worker) {
GCPP_ZONE(ctx, worker, Zones::kGenAttentionQDotK);
if (HWY_LIKELY(last_pos < static_cast<size_t>(div_seq_len.GetDivisor()))) {
// Keep score-buffer indexing unchanged: only KV uses the smaller ring.
const hwy::Divisor div_kv(static_cast<uint32_t>(k.Rows()));
if (HWY_LIKELY(last_pos < k.Rows())) {
// Slightly faster: no wraparound.
for (size_t pos = start_pos; pos <= last_pos; ++pos) {
const float score = Dot(q, k.Row(pos), k.Cols());
Expand All @@ -66,7 +68,7 @@ static HWY_INLINE void QDotK(const size_t start_pos, const size_t last_pos,
} else {
for (size_t pos = start_pos; pos <= last_pos; ++pos) {
const size_t pos_modulo = div_seq_len.Remainder(pos);
const float score = Dot(q, k.Row(pos_modulo), k.Cols());
const float score = Dot(q, k.Row(div_kv.Remainder(pos)), k.Cols());
att[pos_modulo] = score;
}
}
Expand Down Expand Up @@ -98,13 +100,14 @@ void PositionalEncodingQK(float* qk, const size_t layer_idx,
// Accumulates the sum of v (from `kv_cache`) * probability (`att`) into
// `att_out`. Equivalent in gemma/modules.py:
// encoded = jnp.einsum('BTNS,BSNH->BTNH', probs, value_proj)
// `v` is a strided view of the kv cache with dimensions [seq_len, qkv_dim].
// `v` is a strided view of the kv cache with dimensions [cache_rows, qkv_dim].
static HWY_INLINE void WeightedSumV(
const size_t start_pos, const size_t last_pos,
const hwy::Divisor& div_seq_len, const float* HWY_RESTRICT att,
const MatPtrT<KV_t>& v, float* HWY_RESTRICT att_out, ThreadingContext& ctx,
const size_t worker) {
if (HWY_LIKELY(last_pos < static_cast<size_t>(div_seq_len.GetDivisor()))) {
const hwy::Divisor div_kv(static_cast<uint32_t>(v.Rows()));
if (HWY_LIKELY(last_pos < v.Rows())) {
// Slightly faster: no wraparound. Could be replaced with MatMul(att, v) if
// we supported non-transposed B.
// TODO: 2..4x unroll
Expand All @@ -116,12 +119,13 @@ static HWY_INLINE void WeightedSumV(
} else {
{
const size_t pos_mod = div_seq_len.Remainder(start_pos);
MulByConstTo(att[pos_mod], v.Row(pos_mod), att_out, v.Cols(), ctx,
worker);
MulByConstTo(att[pos_mod], v.Row(div_kv.Remainder(start_pos)), att_out,
v.Cols(), ctx, worker);
}
for (size_t pos = start_pos + 1; pos <= last_pos; ++pos) {
const size_t pos_mod = div_seq_len.Remainder(pos);
MulByConstAndAdd(att[pos_mod], v.Row(pos_mod), att_out, v.Cols());
MulByConstAndAdd(att[pos_mod], v.Row(div_kv.Remainder(pos)), att_out,
v.Cols());
}
}
}
Expand Down Expand Up @@ -183,7 +187,6 @@ void DotSoftmaxWeightedSum(const size_t num_tokens, const size_t layer_idx,
// heads that share the same key and value heads.
const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads;

const size_t cache_layer_size = layer_config.CacheLayerSize();
const size_t seq_len =
static_cast<size_t>(activations.div_seq_len.GetDivisor());
// All layers should have the same number of heads.
Expand All @@ -197,7 +200,7 @@ void DotSoftmaxWeightedSum(const size_t num_tokens, const size_t layer_idx,

const size_t qi = div_qbatch.Remainder(tq_idx);
const size_t batch_idx = div_qbatch.Divide(tq_idx);
auto& kv_cache = qbatch.KV(qi).kv_cache;
auto& kv_cache = qbatch.KV(qi).LayerCache(layer_idx);

// Find the token position in the query and calculate
// the range of cache positions to attend to.
Expand All @@ -218,10 +221,10 @@ void DotSoftmaxWeightedSum(const size_t num_tokens, const size_t layer_idx,
// Make strided read-only views into the kv cache for
// this query and head.
const size_t head_offset = (head / kHeadGroups) * qkv_dim * 2;
const size_t kv_head_offset = layer_idx * cache_layer_size + head_offset;
MatPtrT<KV_t> k("k_view", Extents2D(seq_len, qkv_dim));
const size_t kv_head_offset = head_offset;
MatPtrT<KV_t> k("k_view", Extents2D(kv_cache.Rows(), qkv_dim));
k.SetPtr(kv_cache.Row(0) + kv_head_offset, kv_cache.Stride());
MatPtrT<KV_t> v("v_view", Extents2D(seq_len, qkv_dim));
MatPtrT<KV_t> v("v_view", Extents2D(kv_cache.Rows(), qkv_dim));
v.SetPtr(kv_cache.Row(0) + kv_head_offset + qkv_dim, kv_cache.Stride());

SingleDotSoftmaxWeightedSum(pos, start_pos, last_pos, q, k, v, layer_idx,
Expand Down Expand Up @@ -257,7 +260,10 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
const LayerConfig& layer_config = layer.layer_config;
const size_t qkv_dim = layer_config.qkv_dim;
const size_t kv_heads = layer_config.kv_heads;
const size_t cache_layer_size = layer_config.CacheLayerSize();

for (size_t qi = 0; qi < qbatch.Size(); ++qi) {
qbatch.KV(qi).PrepareLayer(layer_idx, num_tokens, qbatch.Pos(qi));
}

// The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim,
// model_dim], which we reshaped to (heads + kv_heads * 2) * qkv_dim rows.
Expand All @@ -266,17 +272,15 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,

// Set up MatMul row pointers for writing to KV, which consists of
// `kv_heads` pairs of (k, v) vectors. This safely handles wraparound
// because rows are computed modulo seq_len.
// because rows are computed modulo each layer's capacity.
MatPtrT<KV_t> kv_rows("kv", Extents2D(activations.pre_att_rms_out.Rows(),
layer.qkv_einsum_w2.Rows()));
for (size_t interleaved_idx = 0; interleaved_idx < num_interleaved;
++interleaved_idx) {
const size_t qi = div_qbatch.Remainder(interleaved_idx);
const size_t batch_idx = div_qbatch.Divide(interleaved_idx);
const size_t cache_pos =
activations.div_seq_len.Remainder(qbatch.Pos(qi) + batch_idx);
env.row_ptrs[0][interleaved_idx] = reinterpret_cast<uint8_t*>(
qbatch.KV(qi).kv_cache.Row(cache_pos) + layer_idx * cache_layer_size);
qbatch.KV(qi).Row(layer_idx, qbatch.Pos(qi) + batch_idx));
}
kv_rows.AttachRowPtrs(env.row_ptrs[0].get());
CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w2,
Expand All @@ -294,11 +298,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
const size_t qi = div_qbatch.Remainder(interleaved_idx);
const size_t batch_idx = div_qbatch.Divide(interleaved_idx);
const size_t pos = qbatch.Pos(qi) + batch_idx;
const size_t cache_pos = activations.div_seq_len.Remainder(pos);
auto& kv_cache = qbatch.KV(qi).kv_cache;
KV_t* HWY_RESTRICT kv = kv_cache.Row(cache_pos) +
layer_idx * cache_layer_size +
head * qkv_dim * 2;
KV_t* HWY_RESTRICT kv =
qbatch.KV(qi).Row(layer_idx, pos) + head * qkv_dim * 2;

HWY_ALIGN float kv_f32[2 * kMaxQKVDim];
const hn::ScalableTag<float> df;
Expand Down
26 changes: 13 additions & 13 deletions gemma/flash_attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,8 @@ void SingleFlashAttention(const size_t start_pos, const size_t last_pos,
float* HWY_RESTRICT att_out, ThreadingContext& ctx,
const size_t worker) {
GCPP_ZONE(ctx, worker, Zones::kFlashAttentionSingleFlashAttention);
const size_t pos_mod = activations.div_seq_len.Remainder(start_pos);
const hwy::Divisor div_kv(static_cast<uint32_t>(k.Rows()));
const size_t pos_mod = div_kv.Remainder(start_pos);
float m = Dot(q, k.Row(pos_mod), k.Cols());
if (float cap = activations.config.att_cap; cap > 0.0f) {
// Compute tanh(x / cap) * cap, being LogitsSoftCap on the scalar x.
Expand All @@ -169,7 +170,7 @@ void SingleFlashAttention(const size_t start_pos, const size_t last_pos,
// This is just a copy of the first token.
MulByConstTo(d, v.Row(pos_mod), att_out, v.Cols(), ctx, worker);
for (size_t pos = start_pos + 1; pos <= last_pos; ++pos) {
const size_t pos_mod = activations.div_seq_len.Remainder(pos);
const size_t pos_mod = div_kv.Remainder(pos);
float x = Dot(q, k.Row(pos_mod), k.Cols());
SingleFlashAttentionStep(x, activations.config.att_cap, m, d,
v.Row(pos_mod), v.Cols(), att_out);
Expand Down Expand Up @@ -275,6 +276,7 @@ void TileFlashAttention(
MatPtrT<float>& att_out, const uint32_t* HWY_RESTRICT out_offsets,
ThreadingContext& ctx, const size_t worker) {
GCPP_ZONE(ctx, worker, Zones::kFlashAttentionTileFlashAttention);
const hwy::Divisor div_kv(static_cast<uint32_t>(k.Rows()));
constexpr int kHTileSize = kNFx8HTileSize;
using DF = hn::ScalableTag<float>;
const DF df;
Expand All @@ -296,7 +298,7 @@ void TileFlashAttention(
while (position + kHTileSize - 1 <= min_last_pos) {
size_t k_pos[kHTileSize];
for (size_t i = 0; i < kHTileSize; ++i) {
k_pos[i] = activations.div_seq_len.Remainder(position + i);
k_pos[i] = div_kv.Remainder(position + i);
}
VF x0, x1, x2, x3, x4, x5, x6, x7;
QDotKTileFloat(df, qT_row, qT_stride, k, k_pos, x0, x1, x2, x3, x4, x5, x6,
Expand Down Expand Up @@ -343,7 +345,7 @@ void TileFlashAttention(
position += kHTileSize;
}
while (position <= max_last_pos) {
size_t k_pos = activations.div_seq_len.Remainder(position);
size_t k_pos = div_kv.Remainder(position);
VF x0 = QDotKVector(df, q_offsets, k_pos, q, k);
if (activations.config.att_cap > 0.0f) {
// Compute tanh(x / cap) * cap, being LogitsSoftCap on the vector.
Expand Down Expand Up @@ -428,6 +430,7 @@ void TileFlashAttention4(
MatPtrT<float>& att_out, const uint32_t* HWY_RESTRICT out_offsets,
ThreadingContext& ctx, const size_t worker) {
GCPP_ZONE(ctx, worker, Zones::kFlashAttentionTileFlashAttention4);
const hwy::Divisor div_kv(static_cast<uint32_t>(k.Rows()));
using DF = hn::ScalableTag<float>;
const DF df;
using VF = hn::Vec<DF>;
Expand All @@ -453,7 +456,7 @@ void TileFlashAttention4(
int32_t k_offsets[kMaxNF];
size_t v_pos[kMaxNF];
for (size_t i = 0; i < kHTileSize; ++i) {
v_pos[i] = activations.div_seq_len.Remainder(position + i);
v_pos[i] = div_kv.Remainder(position + i);
k_offsets[i] = k.Row(v_pos[i]) - k.Row(0);
}
VF x0, x1, x2, x3;
Expand All @@ -476,7 +479,7 @@ void TileFlashAttention4(
position += kHTileSize;
}
while (position <= max_last_pos) {
size_t k_pos = activations.div_seq_len.Remainder(position);
size_t k_pos = div_kv.Remainder(position);
if (position <= last_pos[0]) {
// Past the last position, x0 doesn't count.
float x0 = Dot(q.Row(0) + q_offsets[0], k.Row(k_pos), k.Cols());
Expand Down Expand Up @@ -602,9 +605,6 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism,
// A "head group" in the context of GQA refers to a collection of query
// heads that share the same key and value heads.
const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads;
const size_t cache_layer_size = layer_config.CacheLayerSize();
const size_t seq_len =
static_cast<size_t>(activations.div_seq_len.GetDivisor());
const size_t token_batch = num_tokens * div_qbatch.GetDivisor();
const size_t total_tasks = token_batch * layer_config.heads;

Expand Down Expand Up @@ -696,7 +696,7 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism,
activations.att_out.Row(0);
const size_t kv_index = head / kHeadGroups;
const size_t head_offset = kv_index * qkv_dim * 2;
kv_offsets[offset] = layer_idx * cache_layer_size + head_offset;
kv_offsets[offset] = head_offset;
// If any of the parameters in this if statement differ within this task,
// then we can't use TileFlashAttention. TileFlashAttention requires that
// all rows in the tile have the same K and V matrices, and Q starts at
Expand All @@ -709,10 +709,10 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism,
}
for (size_t offset = 0;
offset < kVTileSize && first_task + offset < total_tasks; ++offset) {
auto& kv_cache = qbatch.KV(qi_indices[offset]).kv_cache;
MatPtrT<KV_t> k("k_view", Extents2D(seq_len, qkv_dim));
auto& kv_cache = qbatch.KV(qi_indices[offset]).LayerCache(layer_idx);
MatPtrT<KV_t> k("k_view", Extents2D(kv_cache.Rows(), qkv_dim));
k.SetPtr(kv_cache.Row(0) + kv_offsets[offset], kv_cache.Stride());
MatPtrT<KV_t> v("v_view", Extents2D(seq_len, qkv_dim));
MatPtrT<KV_t> v("v_view", Extents2D(kv_cache.Rows(), qkv_dim));
v.SetPtr(kv_cache.Row(0) + kv_offsets[offset] + qkv_dim,
kv_cache.Stride());
if (use_tile_attention) {
Expand Down
79 changes: 78 additions & 1 deletion gemma/flash_attention_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ void TestFlashAttention(size_t target_parallelism) {
const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads;
const size_t seq_len =
static_cast<size_t>(attention.div_seq_len.GetDivisor());
auto& kvc = qbatch.KV(0).kv_cache;
auto& kvc = qbatch.KV(0).LayerCache(0);
for (size_t h = 0; h < layer_config.heads; ++h) {
// Make strided views into the kv cache for
// this query and head.
Expand Down Expand Up @@ -164,6 +164,82 @@ void TestFlashAttention(size_t target_parallelism) {
ctx.profiler.PrintResults();
}

// Compare compact rings against full-history buffers with exactly the same
// FP32 arithmetic. Exercise local/global wraparound, runtime growth, prefix-LM,
// and all three FlashAttention tile choices as well as the old attention path.
void TestWindowedKVAttention() {
ThreadingArgs threading;
threading.max_threads = 2;
ThreadingContext ctx(threading);
ModelConfig config(Model::GEMMA3_1B, Type::kF32, PromptWrapping::GEMMA_PT);
config.max_seq_len = 256;
config.num_layers = 1;
config.layer_configs.resize(1);
config.layer_configs[0].heads = 8;
config.attention_window_sizes = {17};
const LayerConfig& layer_config = config.layer_configs[0];
TensorInfoRegistry registry(config);
const LayerWeightsPtrs layer(0, layer_config, registry);
ModelConfig full_config = config;
full_config.attention_window_sizes = {256};
InferenceArgs inference;
inference.seq_len = 256;
inference.prefill_tbatch_size = 8;

for (size_t batch : {1u, 8u, 31u}) {
for (size_t pos : {0u, 23u, 240u, 511u}) {
for (bool prefix : {false, true}) {
for (size_t parallelism : {0u, 1u, 16u, 8192u}) {
SCOPED_TRACE(::testing::Message()
<< "batch=" << batch << " pos=" << pos << " prefix="
<< prefix << " parallelism=" << parallelism);
KVCache compact(config, inference, ctx.allocator);
KVCache full(full_config, inference, ctx.allocator);
compact.PrepareLayer(0, batch, 0);
std::vector<int> tokens(batch, 1);
const size_t end = prefix ? pos + batch : 0;
AllQueries cq(PromptTokens(tokens), pos, end,
hwy::Span<KVCache>(&compact, 1));
AllQueries fq(PromptTokens(tokens), pos, end,
hwy::Span<KVCache>(&full, 1));
QBatch cb(0, 1, cq), fb(0, 1, fq);
std::vector<hwy::AlignedFreeUniquePtr<uint8_t*[]>> cptrs, fptrs;
AttentionActivations ca(config, layer_config, batch, 256,
ctx.allocator, cptrs);
AttentionActivations fa(config, layer_config, batch, 256,
ctx.allocator, fptrs);
for (size_t p = 0; p < pos + batch; ++p) {
for (size_t col = 0; col < full.LayerCache(0).Cols(); ++col) {
const float value =
0.001f * static_cast<float>(
static_cast<int>((p * 13 + col) % 97) - 48);
compact.Row(0, p)[col] = full.Row(0, p)[col] = value;
}
}
SetMat(3, ca.q);
CopyMat(ca.q, fa.q);
for (size_t r = 0; r < batch; ++r) {
// Identical masked score scratch for the old attention path.
std::fill(ca.att.Row(r), ca.att.Row(r) + ca.att.Cols(), -1e30f);
std::fill(fa.att.Row(r), fa.att.Row(r) + fa.att.Cols(), -1e30f);
}
if (parallelism == 0) {
DotSoftmaxWeightedSum(batch, 0, layer, ca, cb, ctx);
DotSoftmaxWeightedSum(batch, 0, layer, fa, fb, ctx);
} else {
FlashAttention(batch, parallelism, 0, layer, ca, cb, ctx);
FlashAttention(batch, parallelism, 0, layer, fa, fb, ctx);
}
for (size_t r = 0; r < batch; ++r) {
ASSERT_EQ(0, std::memcmp(ca.att_out.Row(r), fa.att_out.Row(r),
ca.att_out.Cols() * sizeof(float)));
}
}
}
}
}
}

void TestAttention() {
TestFlashAttention(8192);
TestFlashAttention(2048);
Expand All @@ -180,6 +256,7 @@ HWY_AFTER_NAMESPACE();
namespace gcpp {
HWY_BEFORE_TEST(FlashAttentionTest);
HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestAttention);
HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestWindowedKVAttention);
HWY_AFTER_TEST();

} // namespace gcpp
Expand Down
Loading