Skip to content

Reduce default-flash RAM by reusing BF16 projection scratch #1030

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
Mikyx-1 wants to merge 1 commit into
base: dev
Choose a base branch
from
+121 −23
Open

Reduce default-flash RAM by reusing BF16 projection scratch #1030

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
24 changes: 10 additions & 14 deletions gemma/attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,17 +197,21 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
? static_cast<size_t>(layer_config.kv_share_layer_idx)
: layer_idx;
const bool skip_kv = (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV);
const size_t cache_layer_size = activations.config.layer_configs[kv_layer_idx].CacheLayerSize();

// 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.
CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w1,
/*add=*/nullptr, env, activations.q);

if (skip_kv) return;
// 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.
// Each query reuses one layer of BF16 projections for the current batch.
// Grow for callers that override the configured prefill batch size.
for (size_t qi = 0; qi < qbatch.Size(); ++qi) {
auto& kv = qbatch.KV(qi);
kv.cache->EnsureProjectionRows(num_tokens, layer.qkv_einsum_w2.Rows());
kv.kv_cache = kv.cache->kv_cache;
}
// Set up MatMul row pointers for `kv_heads` pairs of (k, v) vectors.
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;
Expand All @@ -219,12 +223,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
// --seq_len must be large enough to avoid wraparound.
HWY_DASSERT(cache_pos < activations.SeqLen());

const size_t layer_offset = qbatch.KV(qi).cache->layer_flat_offsets.empty()
? kv_layer_idx * cache_layer_size
: qbatch.KV(qi).cache->layer_flat_offsets[kv_layer_idx];

env.row_ptrs[0][interleaved_idx] = reinterpret_cast<uint8_t*>(
qbatch.KV(qi).kv_cache.Row(cache_pos) + layer_offset);
qbatch.KV(qi).kv_cache.Row(token_idx));
}
kv_rows.AttachRowPtrs(env.row_ptrs[0].get());
CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w2,
Expand Down Expand Up @@ -279,11 +279,7 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
// --seq_len must be large enough to avoid wraparound.
HWY_DASSERT(cache_pos < activations.SeqLen());
auto& kv_cache = qbatch.KV(qi).kv_cache;
const size_t layer_offset = qbatch.KV(qi).cache->layer_flat_offsets.empty()
? kv_layer_idx * cache_layer_size
: qbatch.KV(qi).cache->layer_flat_offsets[kv_layer_idx];
KV_t* HWY_RESTRICT kv = kv_cache.Row(cache_pos) +
layer_offset +
KV_t* HWY_RESTRICT kv = kv_cache.Row(token_idx) +
head * qkv_dim * 2;
// Note that k_cache and v_cache are different shapes.
// The innermost dimension of k is 2 values from qkv_dim because they
Expand Down
21 changes: 15 additions & 6 deletions gemma/attention_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ struct TestModelState {
struct TestAttentionState {
TestAttentionState(TestState& state, TestModelState& model_state,
size_t num_tokens, size_t qbatch_size,
AttentionImpl attention_impl)
AttentionImpl attention_impl, bool projection_scratch)
: num_tokens(num_tokens),
qbatch_size(qbatch_size),
batch_size(qbatch_size * num_tokens),
Expand All @@ -112,8 +112,13 @@ struct TestAttentionState {
row_ptrs_),
attention(model_state.config, num_tokens, attention_storage_) {
for (size_t i = 0; i < qbatch_size; ++i) {
kv_caches.emplace_back(model_state.config, inference_args,
state.ctx.allocator);
if (projection_scratch) {
kv_caches.emplace_back(model_state.config, inference_args,
runtime_config, state.ctx.allocator);
} else {
kv_caches.emplace_back(model_state.config, inference_args,
state.ctx.allocator);
}
}
activations.emplace(
runtime_config, model_state.config, runtime_config.prefill_tbatch_size,
Expand Down Expand Up @@ -550,11 +555,12 @@ const float kGoldenQ[kNumTokens][kQBatchSize][kDimsToCompare] = {
0.484799922, 0.0824087635}},
};

void RunAttentionTest(AttentionImpl attention_impl) {
void RunAttentionTest(AttentionImpl attention_impl, bool projection_scratch) {
TestState state;
TestModelState model_state(state);
TestAttentionState attention_state(state, model_state, kNumTokens,
kQBatchSize, attention_impl);
kQBatchSize, attention_impl,
projection_scratch);

GemmaAttention(attention_state.tokens.size(), 0, model_state.layer,
attention_state.attention, *attention_state.qbatch, state.env,
Expand All @@ -569,7 +575,10 @@ void RunAttentionTest(AttentionImpl attention_impl) {
/*q_head=*/0, kGoldenQ);
}

void TestGemmaAttentionFlash() { RunAttentionTest(AttentionImpl::kFlash); }
void TestGemmaAttentionFlash() {
RunAttentionTest(AttentionImpl::kFlash, false);
RunAttentionTest(AttentionImpl::kFlash, true);
}

} // namespace HWY_NAMESPACE
} // namespace gcpp
Expand Down
40 changes: 38 additions & 2 deletions gemma/kv_cache.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,9 +275,20 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
runtime_config.attention_impl == AttentionImpl::kInt8MatrixAccumulation
) {
// clang-format on
kv_is_scratch = runtime_config.attention_impl == AttentionImpl::kFlash &&
!config.is_encoder_decoder && !config.HasMLA();
size_t projection_cols = config.KVCacheCols();
size_t projection_rows = CappedSeqLen(config, inference_args);
if (kv_is_scratch) {
projection_cols = 0;
for (const auto& layer : kv_layer_configs) {
projection_cols = std::max(projection_cols, layer.CacheLayerSize());
}
projection_rows = std::min(
projection_rows, std::max(size_t{1}, runtime_config.prefill_tbatch_size));
}
kv_cache = MatStorageT<KV_t>(
"kv",
Extents2D(CappedSeqLen(config, inference_args), config.KVCacheCols()),
"kv", Extents2D(projection_rows, projection_cols),
allocator, MatPadding::kOdd);
k_cache = MatStorageT<KV_t>(
"k",
Expand All @@ -294,6 +305,8 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
const size_t num_tiles =
hwy::DivCeil(CappedSeqLen(config, inference_args), kTileSize);
tiled_seq_len = num_tiles * kTileSize;
// Default flash reads k_cache/v_cache, never the compact tiled buffers.
if (kv_is_scratch) return;
Type kv_cache_type;
if (runtime_config.attention_impl ==
AttentionImpl::kFlashMatrixAccumulation) {
Expand Down Expand Up @@ -467,10 +480,33 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
InitDSState(config, allocator, ds_state, ds_state_snapshot, ds_state_offsets);
}

void KVCache::EnsureProjectionRows(size_t num_tokens, size_t cols) {
// The constructor without RuntimeConfig also serves legacy callers. Convert
// its projection storage only once we know GemmaAttention is consuming it.
if (!kv_is_scratch || num_tokens > kv_cache.Rows() || cols > kv_cache.Cols()) {
const size_t rows = kv_is_scratch ? std::max(num_tokens, kv_cache.Rows())
: num_tokens;
cols = kv_is_scratch ? std::max(cols, kv_cache.Cols()) : cols;
kv_cache = MatStorageT<KV_t>(
"kv", Extents2D(rows, cols), allocator_, MatPadding::kOdd);
kv_is_scratch = true;
}
}

KVCache KVCache::Copy() {
KVCache copy(kv_cache.Extents(), num_layers, kv_heads, qkv_dim, allocator_);

CopyMat(kv_cache, copy.kv_cache);
copy.kv_is_scratch = kv_is_scratch;
if (kv_is_scratch) {
copy.k_cache = MatStorageT<KV_t>("k", k_cache.Extents(), allocator_,
MatPadding::kPacked);
copy.v_cache = MatStorageT<KV_t>("v", v_cache.Extents(), allocator_,
MatPadding::kPacked);
CopyMat(k_cache, copy.k_cache);
CopyMat(v_cache, copy.v_cache);
copy.k_v_cols = k_v_cols;
}
if (compact_local_kv_cache_ptr.HasPtr()) {
CopyMat(compact_local_kv_cache_ptr, copy.compact_local_kv_cache_ptr);
}
Expand Down
5 changes: 5 additions & 0 deletions gemma/kv_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,11 @@ struct KVCache {
// copy ctor to make the cost explicit.
KVCache Copy();

// Default flash only needs the current layer's projections until transpose.
// Other backends can still use kv_cache as persistent inference state.
bool kv_is_scratch = false;
void EnsureProjectionRows(size_t num_tokens, size_t cols);

size_t SeqLen() const {
if (IsTiled()) {
return tiled_seq_len.value();
Expand Down
54 changes: 53 additions & 1 deletion gemma/kv_cache_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,13 @@ TEST(KVCacheTest, EncoderDecoderUsesDecoderLayerConfig) {
EXPECT_EQ(cache.kv_heads, model_config.decoder_layer_configs[0].kv_heads);
EXPECT_EQ(cache.qkv_dim, model_config.decoder_layer_configs[0].qkv_dim);
EXPECT_EQ(cache.kv_cache.Cols(), model_config.KVCacheCols());
RuntimeConfig runtime_config;
runtime_config.attention_impl = AttentionImpl::kFlash;
KVCache runtime_cache(model_config, inference_args, runtime_config,
ctx.allocator);
EXPECT_FALSE(runtime_cache.kv_is_scratch);
EXPECT_EQ(runtime_cache.kv_cache.Rows(), inference_args.seq_len);
EXPECT_EQ(runtime_cache.kv_cache.Cols(), model_config.KVCacheCols());
}

// Layers that reuse an earlier layer's K/V own no region of the cache.
Expand All @@ -79,7 +86,52 @@ TEST(KVCacheTest, SharedLayersReserveNoCache) {
EXPECT_EQ(cache.layer_flat_offsets[15], cache.layer_flat_offsets[13]);
EXPECT_EQ(cache.layer_k_v_offsets[15], cache.layer_k_v_offsets[13]);
EXPECT_EQ(cache.layer_kv_head_offsets[15], cache.layer_kv_head_offsets[13]);
EXPECT_EQ(cache.kv_cache.Cols(), model_config.KVCacheCols());
EXPECT_TRUE(cache.kv_is_scratch);
// Global layers project 1024 values; local layers project only 512.
EXPECT_EQ(cache.kv_cache.Cols(), 1024);
}

TEST(KVCacheTest, FlashScratchPreservesSequenceAndHistory) {
ModelConfig config(Model::GEMMA3_270M, Type::kSFP, PromptWrapping::GEMMA_IT);
config.num_layers = 2;
config.layer_configs.resize(2);
config.attention_window_sizes.resize(2);
InferenceArgs args;
args.seq_len = 32768;
RuntimeConfig runtime;
runtime.attention_impl = AttentionImpl::kFlash;
runtime.prefill_tbatch_size = 16;
ThreadingArgs threading;
ThreadingContext ctx(threading);
KVCache cache(config, args, runtime, ctx.allocator);
ASSERT_TRUE(cache.kv_is_scratch);
EXPECT_EQ(cache.SeqLen(), 32768);
EXPECT_EQ(cache.ToPtr().SeqLen(), 32768);
EXPECT_FALSE(cache.ToPtr().IsEmpty());
EXPECT_EQ(cache.kv_cache.Rows(), 16);
EXPECT_EQ(cache.kv_cache.Cols(), config.layer_configs[0].CacheLayerSize());
EXPECT_FALSE(cache.compact_kv_cache_ptr.HasPtr());

ZeroInit(cache.kv_cache);
ZeroInit(cache.k_cache);
ZeroInit(cache.v_cache);
cache.k_cache.Row(32767)[0] = hwy::BF16FromF32(3.0f);
cache.v_cache.Row(32767)[0] = hwy::BF16FromF32(5.0f);
cache.EnsureProjectionRows(33, cache.kv_cache.Cols());
EXPECT_EQ(cache.kv_cache.Rows(), 33);
cache.EnsureProjectionRows(1, cache.kv_cache.Cols());
EXPECT_EQ(cache.kv_cache.Rows(), 33);
const size_t wider_cols = 2 * cache.kv_cache.Cols();
cache.EnsureProjectionRows(1, wider_cols);
EXPECT_EQ(cache.kv_cache.Rows(), 33);
EXPECT_EQ(cache.kv_cache.Cols(), wider_cols);
ZeroInit(cache.kv_cache);
KVCache copy = cache.Copy();
EXPECT_TRUE(copy.kv_is_scratch);
EXPECT_EQ(copy.SeqLen(), 32768);
EXPECT_EQ(hwy::F32FromBF16(copy.k_cache.Row(32767)[0]), 3.0f);
EXPECT_EQ(hwy::F32FromBF16(copy.v_cache.Row(32767)[0]), 5.0f);
EXPECT_NE(copy.k_cache.Row(0), cache.k_cache.Row(0));
}

} // namespace
Expand Down