Skip to content

Extend CompressAndTransposeQueriesMatrixAccumulationImpl and #1007

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

Merged
copybara-service merged 1 commit into from
Sep 3, 2026
+185 −51
Merged

Extend CompressAndTransposeQueriesMatrixAccumulationImpl and #1007

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
157 changes: 106 additions & 51 deletions gemma/tiled_attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,24 +73,26 @@ static HWY_INLINE void MergeOnlineSoftmax(
}

template <typename T>
T AbsMaxOfSpan(hwy::Span<const T> span) {
hn::ScalableTag<T> dt;
using VT = hn::Vec<decltype(dt)>;
VT max_vec = hn::Set(dt, 0.0f);
const size_t lanes = hn::Lanes(dt);
float AbsMaxOfSpan(hwy::Span<const T> span) {
namespace hn = hwy::HWY_NAMESPACE;
const hn::ScalableTag<float> df;
using VF = hn::Vec<decltype(df)>;
VF max_vec = hn::Zero(df);
HWY_LANES_CONSTEXPR size_t N = hn::Lanes(df);
HWY_LANES_CONSTEXPR size_t step = 2 * N;
const PackedSpan<const T> packed_span(span.data(), span.size());
size_t i = 0;
// Process full vectors using LoadU.
for (; i + lanes <= span.size(); i += lanes) {
const VT vec = hn::Abs(hn::LoadU(dt, span.data() + i));
max_vec = hn::Max(max_vec, vec);
for (; i + step <= span.size(); i += step) {
VF v0, v1;
Decompress2(df, packed_span, i, v0, v1);
max_vec = hn::Max(max_vec, hn::Max(hn::Abs(v0), hn::Abs(v1)));
}
// Process remaining elements using LoadN.
const size_t remaining = span.size() - i;
if (HWY_UNLIKELY(remaining != 0)) {
const VT vec = hn::Abs(hn::LoadN(dt, span.data() + i, remaining));
max_vec = hn::Max(max_vec, vec);
float max_scalar = 0.0f;
for (; i < span.size(); ++i) {
max_scalar = HWY_MAX(max_scalar,
hwy::ScalarAbs(hwy::ConvertScalarTo<float>(span[i])));
}
return hn::ReduceMax(dt, max_vec);
return HWY_MAX(hn::ReduceMax(df, max_vec), max_scalar);
}

// Forked from ComputeQKV. But it stores the K/V in the tiled format
Expand Down Expand Up @@ -561,37 +563,53 @@ template <typename QueryProvider>
HWY_INLINE void CompressAndTransposeQueriesMatrixAccumulationImpl(
QueryProvider query_provider, BF16* packed_queries, size_t num_queries,
size_t qkv_dim) {
HWY_DASSERT(qkv_dim % 4 == 0);

namespace hn = hwy::HWY_NAMESPACE;
const hn::Full128<float> df;
using InT = hwy::RemoveCvRef<hwy::RemovePtr<decltype(query_provider(0))>>;
const hn::Full128<BF16> dbf16;
constexpr size_t kL = 4;
const hn::Half<decltype(dbf16)> dbf_half;
const hn::Full128<float> df;
const size_t kL = hn::Lanes(dbf_half);

using V_BF16 = hn::Vec<decltype(dbf16)>;
using V_BF16_Half = hn::Vec<decltype(dbf_half)>;
using V_F32 = hn::Vec<decltype(df)>;

HWY_DASSERT(qkv_dim % kL == 0);

auto pack4x2 = [&](const InT* q0,
const InT* q1) HWY_ATTR -> V_BF16 {
if constexpr (IsBF16<InT>()) {
const V_BF16_Half v0 = hn::LoadU(dbf_half, q0);
const V_BF16_Half v1 =
q1 != nullptr ? hn::LoadU(dbf_half, q1) : hn::Zero(dbf_half);
return hn::Combine(dbf16, v1, v0);
} else {
const V_F32 v0 = hn::LoadU(df, q0);
const V_F32 v1 = q1 != nullptr ? hn::LoadU(df, q1) : hn::Zero(df);
return hn::OrderedDemote2To(dbf16, v0, v1);
}
};

auto pack_pair = [&](const InT* q0, const InT* q1, BF16* out) HWY_ATTR {
if (q1 != nullptr) {
for (size_t d = 0; d < qkv_dim; d += kL) {
hn::StoreU(pack4x2(q0 + d, q1 + d), dbf16, out + d * 2);
}
} else {
for (size_t d = 0; d < qkv_dim; d += kL) {
hn::StoreU(pack4x2(q0 + d, nullptr), dbf16, out + d * 2);
}
}
};

size_t p = 0;
for (; p < num_queries / 2; ++p) {
const float* q0 = query_provider(2 * p);
const float* q1 = query_provider(2 * p + 1);
BF16* out = packed_queries + 2 * p * qkv_dim;

for (size_t d = 0; d < qkv_dim; d += kL) {
auto v0 = hn::LoadU(df, q0 + d);
auto v1 = hn::LoadU(df, q1 + d);
auto A = hn::OrderedDemote2To(dbf16, v0, v1);
hn::StoreU(A, dbf16, out + d * 2);
}
pack_pair(query_provider(2 * p), query_provider(2 * p + 1),
packed_queries + 2 * p * qkv_dim);
}

if (num_queries % 2 != 0) {
const float* q0 = query_provider(2 * p);
BF16* out = packed_queries + 2 * p * qkv_dim;
auto zero = hn::Zero(df);

for (size_t d = 0; d < qkv_dim; d += kL) {
auto v0 = hn::LoadU(df, q0 + d);
auto A = hn::OrderedDemote2To(dbf16, v0, zero);
hn::StoreU(A, dbf16, out + d * 2);
}
pack_pair(query_provider(2 * p), nullptr,
packed_queries + 2 * p * qkv_dim);
}
}

Expand All @@ -604,20 +622,36 @@ void CompressAndTransposeQueriesMatrixAccumulation(const float* raw_queries,
num_queries, qkv_dim);
}

void CompressAndTransposeQueriesMatrixAccumulationFromBF16(
const BF16* raw_queries, BF16* packed_queries, size_t num_queries,
size_t qkv_dim) {
CompressAndTransposeQueriesMatrixAccumulationImpl(
[&](size_t idx) { return raw_queries + idx * qkv_dim; }, packed_queries,
num_queries, qkv_dim);
}

void CompressAndTransposeQueriesMatrixAccumulationNonContiguous(
hwy::Span<const float* const> input, BF16* packed_queries, size_t qkv_dim) {
CompressAndTransposeQueriesMatrixAccumulationImpl(
[&](size_t idx) { return input[idx]; }, packed_queries, input.size(),
qkv_dim);
}

void CompressAndTransposeQueriesMatrixAccumulationNonContiguousFromBF16(
hwy::Span<const BF16* const> input, BF16* packed_queries, size_t qkv_dim) {
CompressAndTransposeQueriesMatrixAccumulationImpl(
[&](size_t idx) { return input[idx]; }, packed_queries, input.size(),
qkv_dim);
}

template <typename QueryProvider>
HWY_INLINE void CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(
QueryProvider query_provider, int8_t* HWY_RESTRICT packed_queries,
float* HWY_RESTRICT packed_scales, size_t num_queries, size_t qkv_dim) {
HWY_DASSERT(qkv_dim % 8 == 0);

namespace hn = hwy::HWY_NAMESPACE;
using InT = hwy::RemoveCvRef<hwy::RemovePtr<decltype(query_provider(0))>>;
const hn::Full128<float> df;
const hn::Full128<int16_t> di16;
const hn::Full128<int8_t> di8;
Expand All @@ -629,15 +663,18 @@ HWY_INLINE void CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(

size_t p = 0;
for (; p < num_queries / 2; ++p) {
const float* q0 = query_provider(2 * p);
const float* q1 = query_provider(2 * p + 1);
const InT* q0 = query_provider(2 * p);
const InT* q1 = query_provider(2 * p + 1);
int8_t* out = packed_queries + 2 * p * qkv_dim;
float* out_scale0 = packed_scales + 2 * p;
float* out_scale1 = packed_scales + (2 * p + 1);

const PackedSpan<const InT> span_q0(q0, qkv_dim);
const PackedSpan<const InT> span_q1(q1, qkv_dim);

// 1. Compute single scale per query over the entire qkv_dim
float max_abs_q0 = AbsMaxOfSpan(hwy::Span<const float>(q0, qkv_dim));
float max_abs_q1 = AbsMaxOfSpan(hwy::Span<const float>(q1, qkv_dim));
float max_abs_q0 = AbsMaxOfSpan(hwy::Span<const InT>(q0, qkv_dim));
float max_abs_q1 = AbsMaxOfSpan(hwy::Span<const InT>(q1, qkv_dim));

float scale0_raw = max_abs_q0 == 0.0f ? 1.0f : max_abs_q0 / 127.0f;
float scale1_raw = max_abs_q1 == 0.0f ? 1.0f : max_abs_q1 / 127.0f;
Expand All @@ -658,15 +695,15 @@ HWY_INLINE void CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(

for (size_t d = 0; d < qkv_dim; d += 8) {
// 2. Load and quantize Q0 (8 channels)
V_F32 q0_L = hn::LoadU(df, q0 + d);
V_F32 q0_H = hn::LoadU(df, q0 + d + 4);
V_F32 q0_L, q0_H;
Decompress2(df, span_q0, d, q0_L, q0_H);
V_I32 q0_L_scaled = hn::NearestInt(hn::Mul(q0_L, inv_scale0));
V_I32 q0_H_scaled = hn::NearestInt(hn::Mul(q0_H, inv_scale0));
V_I16 q0_i16 = hn::OrderedDemote2To(di16, q0_L_scaled, q0_H_scaled);

// 3. Load and quantize Q1 (8 channels)
V_F32 q1_L = hn::LoadU(df, q1 + d);
V_F32 q1_H = hn::LoadU(df, q1 + d + 4);
V_F32 q1_L, q1_H;
Decompress2(df, span_q1, d, q1_L, q1_H);
V_I32 q1_L_scaled = hn::NearestInt(hn::Mul(q1_L, inv_scale1));
V_I32 q1_H_scaled = hn::NearestInt(hn::Mul(q1_H, inv_scale1));
V_I16 q1_i16 = hn::OrderedDemote2To(di16, q1_L_scaled, q1_H_scaled);
Expand All @@ -679,12 +716,14 @@ HWY_INLINE void CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(
}

if (num_queries % 2 != 0) {
const float* q0 = query_provider(2 * p);
const InT* q0 = query_provider(2 * p);
int8_t* out = packed_queries + 2 * p * qkv_dim;
float* out_scale0 = packed_scales + 2 * p;
V_I16 zero_i16 = hn::Zero(di16);

float max_abs_q0 = AbsMaxOfSpan(hwy::Span<const float>(q0, qkv_dim));
const PackedSpan<const InT> span_q0(q0, qkv_dim);

float max_abs_q0 = AbsMaxOfSpan(hwy::Span<const InT>(q0, qkv_dim));

float scale0_raw = max_abs_q0 == 0.0f ? 1.0f : max_abs_q0 / 127.0f;
gcpp::KV_microscale_t scale0_bf16 =
Expand All @@ -696,8 +735,8 @@ HWY_INLINE void CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(
V_F32 inv_scale0 = hn::Set(df, 1.0f / scale0);

for (size_t d = 0; d < qkv_dim; d += 8) {
V_F32 q0_L = hn::LoadU(df, q0 + d);
V_F32 q0_H = hn::LoadU(df, q0 + d + 4);
V_F32 q0_L, q0_H;
Decompress2(df, span_q0, d, q0_L, q0_H);
V_I32 q0_L_scaled = hn::NearestInt(hn::Mul(q0_L, inv_scale0));
V_I32 q0_H_scaled = hn::NearestInt(hn::Mul(q0_H, inv_scale0));
V_I16 q0_i16 = hn::OrderedDemote2To(di16, q0_L_scaled, q0_H_scaled);
Expand All @@ -718,6 +757,14 @@ void CompressAndQuantizeQueriesMatrixAccumulationInt8(const float* raw_queries,
packed_scales, num_queries, qkv_dim);
}

void CompressAndQuantizeQueriesMatrixAccumulationInt8FromBF16(
const BF16* raw_queries, int8_t* packed_queries, float* packed_scales,
size_t num_queries, size_t qkv_dim) {
CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(
[&](size_t idx) { return raw_queries + idx * qkv_dim; }, packed_queries,
packed_scales, num_queries, qkv_dim);
}

void CompressAndQuantizeQueriesMatrixAccumulationInt8NonContiguous(
hwy::Span<const float* const> input, int8_t* packed_queries,
float* packed_scales, size_t qkv_dim) {
Expand All @@ -726,6 +773,14 @@ void CompressAndQuantizeQueriesMatrixAccumulationInt8NonContiguous(
input.size(), qkv_dim);
}

void CompressAndQuantizeQueriesMatrixAccumulationInt8NonContiguousFromBF16(
hwy::Span<const BF16* const> input, int8_t* packed_queries,
float* packed_scales, size_t qkv_dim) {
CompressAndQuantizeQueriesMatrixAccumulationInt8Impl(
[&](size_t idx) { return input[idx]; }, packed_queries, packed_scales,
input.size(), qkv_dim);
}

// clang-format off
// Schedules TiledFlashAttention for all heads, tokens and batch.
// Returns partial results in the same order as queries in `activations.q`.
Expand Down
6 changes: 6 additions & 0 deletions gemma/tiled_attention.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,15 @@ namespace gcpp {
BF16* packed_queries, \
size_t num_queries, \
size_t qkv_dim); \
void CompressAndTransposeQueriesMatrixAccumulationFromBF16( \
const BF16* raw_queries, BF16* packed_queries, size_t num_queries, \
size_t qkv_dim); \
void CompressAndQuantizeQueriesMatrixAccumulationInt8( \
const float* raw_queries, int8_t* packed_queries, float* packed_scales, \
size_t num_queries, size_t qkv_dim); \
void CompressAndQuantizeQueriesMatrixAccumulationInt8FromBF16( \
const BF16* raw_queries, int8_t* packed_queries, float* packed_scales, \
size_t num_queries, size_t qkv_dim); \
/* NOLINTNEXTLINE(google-readability-namespace-comments) */ \
} // namespace NAMESPACE

Expand Down
73 changes: 73 additions & 0 deletions gemma/tiled_attention_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,77 @@ void TestCompressQueries() {
}
}

void TestCompressQueriesMatrixAccumulationBF16() {
ThreadingArgs threading_args;
ThreadingContext ctx(threading_args);

for (size_t qkv_dim : {24, 40, 64}) {
for (size_t num_queries : {5, 6, 24, 25}) {
const size_t num_queries_rounded = hwy::RoundUpTo(num_queries, 2);
AlignedPtr<float[]> input_f32 =
ctx.allocator.Alloc<float>(qkv_dim * num_queries);
AlignedPtr<BF16[]> input_bf16 =
ctx.allocator.Alloc<BF16>(qkv_dim * num_queries);

for (size_t i = 0; i < num_queries; ++i) {
for (size_t j = 0; j < qkv_dim; ++j) {
float val = 0.01f * (i + 1) / (j + 1);
BF16 val_bf16 = hwy::ConvertScalarTo<BF16>(val);
input_bf16[i * qkv_dim + j] = val_bf16;
input_f32[i * qkv_dim + j] = hwy::ConvertScalarTo<float>(val_bf16);
}
}

// 1. Test BF16 matrix accumulation compression
AlignedPtr<BF16[]> packed_from_f32 =
ctx.allocator.Alloc<BF16>(qkv_dim * num_queries_rounded);
AlignedPtr<BF16[]> packed_from_bf16 =
ctx.allocator.Alloc<BF16>(qkv_dim * num_queries_rounded);

CompressAndTransposeQueriesMatrixAccumulation(
input_f32.get(), packed_from_f32.get(), num_queries, qkv_dim);
CompressAndTransposeQueriesMatrixAccumulationFromBF16(
input_bf16.get(), packed_from_bf16.get(), num_queries, qkv_dim);

for (size_t i = 0; i < num_queries_rounded * qkv_dim; ++i) {
EXPECT_EQ(hwy::ConvertScalarTo<float>(packed_from_f32[i]),
hwy::ConvertScalarTo<float>(packed_from_bf16[i]))
<< "BF16 matrix accumulation mismatch at index " << i
<< " with num_queries=" << num_queries;
}

// 2. Test Int8 matrix accumulation compression
AlignedPtr<int8_t[]> packed_int8_from_f32 =
ctx.allocator.Alloc<int8_t>(qkv_dim * num_queries_rounded);
AlignedPtr<int8_t[]> packed_int8_from_bf16 =
ctx.allocator.Alloc<int8_t>(qkv_dim * num_queries_rounded);
AlignedPtr<float[]> scales_from_f32 =
ctx.allocator.Alloc<float>(num_queries_rounded);
AlignedPtr<float[]> scales_from_bf16 =
ctx.allocator.Alloc<float>(num_queries_rounded);

CompressAndQuantizeQueriesMatrixAccumulationInt8(
input_f32.get(), packed_int8_from_f32.get(), scales_from_f32.get(),
num_queries, qkv_dim);
CompressAndQuantizeQueriesMatrixAccumulationInt8FromBF16(
input_bf16.get(), packed_int8_from_bf16.get(), scales_from_bf16.get(),
num_queries, qkv_dim);

for (size_t i = 0; i < num_queries; ++i) {
EXPECT_FLOAT_EQ(scales_from_f32[i], scales_from_bf16[i])
<< "Int8 scale mismatch for query " << i
<< " with num_queries=" << num_queries;
}

for (size_t i = 0; i < num_queries_rounded * qkv_dim; ++i) {
EXPECT_EQ(packed_int8_from_f32[i], packed_int8_from_bf16[i])
<< "Int8 packed query mismatch at index " << i
<< " with num_queries=" << num_queries;
}
}
}
}

void TestLocalAttentionForAllHeadsTokensAndBatch() {
size_t qkv_dim = 64;
size_t kv_seq_len = 64;
Expand Down Expand Up @@ -826,6 +897,8 @@ HWY_AFTER_NAMESPACE();
namespace gcpp {
HWY_BEFORE_TEST(TiledAttentionTest);
HWY_EXPORT_AND_TEST_P(TiledAttentionTest, TestCompressQueries);
HWY_EXPORT_AND_TEST_P(TiledAttentionTest,
TestCompressQueriesMatrixAccumulationBF16);
// TODO() Fix the goldens for the change in KV_t to BF16
// HWY_EXPORT_AND_TEST_P(TiledAttentionTest,
// TestLocalAttentionForAllHeadsTokensAndBatch);
Expand Down
Loading