Skip to content

Internal change. #977

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
copybara-service wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Internal change. #977

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
76 changes: 57 additions & 19 deletions compression/q4_0-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,25 +97,57 @@ class Q4_0Codec {
template <class D, typename Raw = hn::TFromD<D>>
static HWY_INLINE void DequantizeBlock(D d, const uint8_t* HWY_RESTRICT block_ptr,
Raw* HWY_RESTRICT raw) {
const hn::Repartition<float, D> df;
const hn::Rebind<int32_t, decltype(df)> di32;
const hn::Rebind<int16_t, decltype(di32)> di16;
const hn::Rebind<int8_t, decltype(di16)> di8;
const hn::Rebind<uint8_t, decltype(di8)> du8;
const hn::Full128<uint8_t> du8_128;
const hn::Full128<int8_t> di8_128;
const hn::Full128<int16_t> di16_128;
const hn::Full128<int32_t> di32_128;
const hn::Full128<float> df_128;

using T = ScaleT;
T scale;
hwy::CopyBytes(block_ptr, &scale, sizeof(T));
const float scale_f = hwy::F32FromBF16(scale);
const auto vd = hn::Set(df, scale_f);
const auto vd_128 = hn::Set(df_128, scale_f);

const uint8_t* qs_ptr = block_ptr + sizeof(T);
const size_t N = hn::Lanes(df);
const size_t num_vectors = 32 / N;

for (size_t v_idx = 0; v_idx < num_vectors; ++v_idx) {
const auto out = DequantizeLanes(df, du8, di8, qs_ptr, 0, v_idx * N, N, vd);
StoreRaw(df, out, raw + v_idx * N);
const auto raw_16 = hn::LoadU(du8_128, qs_ptr);
const auto mask_0f = hn::Set(du8_128, 0x0F);
const auto offset_8 = hn::Set(di8_128, 8);

const auto low_i8 =
hn::Sub(hn::BitCast(di8_128, hn::And(raw_16, mask_0f)), offset_8);
const auto high_i8 =
hn::Sub(hn::BitCast(di8_128, hn::ShiftRight<4>(raw_16)), offset_8);

const auto low_i16_0 = hn::PromoteLowerTo(di16_128, low_i8);
const auto low_i16_1 = hn::PromoteUpperTo(di16_128, low_i8);
const auto high_i16_0 = hn::PromoteLowerTo(di16_128, high_i8);
const auto high_i16_1 = hn::PromoteUpperTo(di16_128, high_i8);

const auto f0 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, low_i16_0)), vd_128);
const auto f1 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, low_i16_0)), vd_128);
const auto f2 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, low_i16_1)), vd_128);
const auto f3 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, low_i16_1)), vd_128);
const auto f4 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, high_i16_0)), vd_128);
const auto f5 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, high_i16_0)), vd_128);
const auto f6 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, high_i16_1)), vd_128);
const auto f7 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, high_i16_1)), vd_128);

if constexpr (hwy::IsSame<Raw, float>()) {
hn::StoreU(f0, df_128, raw + 0);
hn::StoreU(f1, df_128, raw + 4);
hn::StoreU(f2, df_128, raw + 8);
hn::StoreU(f3, df_128, raw + 12);
hn::StoreU(f4, df_128, raw + 16);
hn::StoreU(f5, df_128, raw + 20);
hn::StoreU(f6, df_128, raw + 24);
hn::StoreU(f7, df_128, raw + 28);
} else {
const hn::Full128<hwy::bfloat16_t> dbf_128;
hn::StoreU(hn::OrderedDemote2To(dbf_128, f0, f1), dbf_128, raw + 0);
hn::StoreU(hn::OrderedDemote2To(dbf_128, f2, f3), dbf_128, raw + 8);
hn::StoreU(hn::OrderedDemote2To(dbf_128, f4, f5), dbf_128, raw + 16);
hn::StoreU(hn::OrderedDemote2To(dbf_128, f6, f7), dbf_128, raw + 24);
}
}

Expand Down Expand Up @@ -259,18 +291,24 @@ class Q4_0Codec {
HWY_DASSERT(current_packed_ofs % kBlockSize == 0);

const size_t num_full_blocks = num_to_decompress / kBlockSize;
for (size_t b = 0; b < num_full_blocks; ++b) {
const uint8_t* block_ptr =
&packed.ptr->byte + BlockByteOffset(current_packed_ofs);
const uint8_t* block_ptr =
&packed.ptr->byte + BlockByteOffset(current_packed_ofs);
size_t b = 0;
for (; b + 1 < num_full_blocks; b += 2) {
DequantizeBlock(d, block_ptr, current_raw);
DequantizeBlock(d, block_ptr + 18, current_raw + 32);
block_ptr += 36;
current_raw += 64;
}
if (b < num_full_blocks) {
DequantizeBlock(d, block_ptr, current_raw);
current_packed_ofs += kBlockSize;
current_raw += kBlockSize;
block_ptr += 18;
current_raw += 32;
}
current_packed_ofs += num_full_blocks * kBlockSize;

const size_t remaining = num_to_decompress % kBlockSize;
if (remaining != 0) {
const uint8_t* block_ptr =
&packed.ptr->byte + BlockByteOffset(current_packed_ofs);
HWY_ALIGN Raw temp[kBlockSize];
DequantizeBlock(d, block_ptr, temp);
memcpy(current_raw, temp, remaining * sizeof(Raw));
Expand Down
42 changes: 33 additions & 9 deletions compression/test_util-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,9 +119,15 @@ MatStorageT<MatT> GenerateMat(const Extents2D& extents, MatPadding padding,
f = -f; // Also generate some negative values.
row[c] = f;
}
Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
MakeSpan(compressed.Row(r), extents.cols),
/*packed_ofs=*/0);
if constexpr (IsQ4_0Stream<MatT>()) {
Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
compressed.Span(),
/*packed_ofs=*/r * extents.cols);
} else {
Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
MakeSpan(compressed.Row(r), extents.cols),
/*packed_ofs=*/0);
}
});

compressed.SetScale(0.6f); // Arbitrary value, different from 1.
Expand All @@ -148,9 +154,15 @@ MatStorageT<MatT> GenerateTransposedMat(const Extents2D extents,
f = -f; // Also generate some negative values.
row[c] = f;
}
Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
MakeSpan(compressed.Row(r), extents.cols),
/*packed_ofs=*/0);
if constexpr (IsQ4_0Stream<MatT>()) {
Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
compressed.Span(),
/*packed_ofs=*/r * extents.cols);
} else {
Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
MakeSpan(compressed.Row(r), extents.cols),
/*packed_ofs=*/0);
}
});

// Arbitrary value, different from 1, must match `GenerateMat`.
Expand Down Expand Up @@ -205,15 +217,24 @@ void AssertClose(const MatPtrT<TA>& A, const MatPtrT<TB>& B,
MatStorageT<float> c_slow_batch("c_slow_batch", Extents2D(A.Rows(), B_rows),
allocator, MatPadding::kOdd);
for (size_t m = 0; m < A.Rows(); ++m) {
DecompressAndZeroPad(df, MakeSpan(A.Row(m), cols), 0, a_batch.Row(m), cols);
if constexpr (IsQ4_0Stream<TA>()) {
DecompressAndZeroPad(df, A.PaddedSpan(), m * cols, a_batch.Row(m), cols);
} else {
DecompressAndZeroPad(df, MakeSpan(A.Row(m), cols), 0, a_batch.Row(m), cols);
}
DecompressAndZeroPad(df, MakeSpan(C.Row(m), B_rows), 0, c_batch.Row(m),
B_rows);
DecompressAndZeroPad(df, MakeSpan(C_slow.Row(m), B_rows), 0,
c_slow_batch.Row(m), B_rows);
}
for (size_t n = 0; n < B_rows; ++n) {
DecompressAndZeroPad(df, MakeSpan(B.Row(n), cols), 0, b_trans_batch.Row(n),
cols);
if constexpr (IsQ4_0Stream<TB>()) {
DecompressAndZeroPad(df, B.PaddedSpan(), n * cols, b_trans_batch.Row(n),
cols);
} else {
DecompressAndZeroPad(df, MakeSpan(B.Row(n), cols), 0, b_trans_batch.Row(n),
cols);
}
}

// MatMul rounds inputs to BF16, so error is proportional to the max input
Expand All @@ -231,6 +252,9 @@ void AssertClose(const MatPtrT<TA>& A, const MatPtrT<TB>& B,
if (IsF32<TA>() || IsF32<TB>()) {
tolerance += 2 * max_abs * eps_bf16;
}
if constexpr (IsQ4_0Stream<TB>()) {
tolerance += 0.02 * norm;
}

if (tolerance > 500.0) {
HWY_WARN("high tolerance %f norm %f maxabs %f\n", tolerance, norm, max_abs);
Expand Down
99 changes: 89 additions & 10 deletions evals/gemma_batch_bench.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,15 +36,11 @@ GemmaEnv* s_env = nullptr;

class GemmaBatchBench : public ::testing::Test {
protected:
std::vector<std::string> BatchGemmaReply(
QueryResultAndMetrics BatchGemmaReplyWithMetrics(
const std::vector<std::string>& inputs) {
s_env->MutableConfig().temperature = 0.0f; // deterministic
s_env->MutableConfig().verbosity = 2;
std::vector<std::string> replies;
for (const QueryResult& result : s_env->BatchQueryModel(inputs)) {
replies.push_back(result.response);
}
return replies;
return s_env->BatchQueryModelWithMetrics(inputs);
}
};

Expand Down Expand Up @@ -128,16 +124,99 @@ std::vector<std::string> GenerateInputs() {
TEST_F(GemmaBatchBench, RandomQuestionsBatched) {
s_env->SetMaxGeneratedTokens(12);
const std::vector<std::string> inputs = GenerateInputs();
// Run multiple times so that auto-tuning is closer to complete.
for (size_t rep = 0; rep < 4; ++rep) {
std::vector<std::string> responses = BatchGemmaReply(inputs);
constexpr size_t kNumReps = 7;

std::vector<double> prefill_speeds;
std::vector<double> generate_speeds;
size_t total_prefill_tokens = 0;
double total_prefill_duration = 0.0;
size_t total_generate_tokens = 0;
double total_generate_duration = 0.0;

size_t warm_prefill_tokens = 0;
double warm_prefill_duration = 0.0;
size_t warm_generate_tokens = 0;
double warm_generate_duration = 0.0;

for (size_t rep = 0; rep < kNumReps; ++rep) {
QueryResultAndMetrics result = BatchGemmaReplyWithMetrics(inputs);
const std::vector<QueryResult>& responses = result.query_results;
const TimingInfo& timing = result.timing_info;

const double prefill_tok_sec =
timing.prefill_duration > 0.0
? static_cast<double>(timing.prefill_tokens) /
timing.prefill_duration
: 0.0;
const double gen_tok_sec =
timing.generate_duration > 0.0
? static_cast<double>(timing.tokens_generated) /
timing.generate_duration
: 0.0;

prefill_speeds.push_back(prefill_tok_sec);
generate_speeds.push_back(gen_tok_sec);

total_prefill_tokens += timing.prefill_tokens;
total_prefill_duration += timing.prefill_duration;
total_generate_tokens += timing.tokens_generated;
total_generate_duration += timing.generate_duration;

if (rep > 0) {
warm_prefill_tokens += timing.prefill_tokens;
warm_prefill_duration += timing.prefill_duration;
warm_generate_tokens += timing.tokens_generated;
warm_generate_duration += timing.generate_duration;
}

for (size_t i = 0; i < HWY_MIN(hwy::Unpredictable1() * 3, responses.size());
++i) {
fprintf(stderr, "Rep %zu batch answer %zu '%s'\n\n", rep, i,
responses[i].c_str());
responses[i].response.c_str());
}
PROFILER_PRINT_RESULTS();
}

const double avg_prefill =
total_prefill_duration > 0.0
? static_cast<double>(total_prefill_tokens) / total_prefill_duration
: 0.0;
const double avg_generate =
total_generate_duration > 0.0
? static_cast<double>(total_generate_tokens) / total_generate_duration
: 0.0;

const double warm_avg_prefill =
warm_prefill_duration > 0.0
? static_cast<double>(warm_prefill_tokens) / warm_prefill_duration
: 0.0;
const double warm_avg_generate =
warm_generate_duration > 0.0
? static_cast<double>(warm_generate_tokens) / warm_generate_duration
: 0.0;

fprintf(stderr,
"\n============================================================\n");
fprintf(stderr,
"[ Gemma Batch Benchmark Summary (%zu Repetitions) ]\n", kNumReps);
for (size_t rep = 0; rep < kNumReps; ++rep) {
fprintf(stderr,
" Rep %zu: Prefill = %7.2f tok/s | Generate = %7.2f tok/s%s\n",
rep, prefill_speeds[rep], generate_speeds[rep],
rep == 0 ? " (warmup / autotune)" : "");
}
fprintf(stderr,
"------------------------------------------------------------\n");
fprintf(stderr,
"Overall Average: Prefill = %7.2f tok/s | Generate = %7.2f tok/s\n",
avg_prefill, avg_generate);
if (kNumReps > 1) {
fprintf(stderr,
"Warm Average : Prefill = %7.2f tok/s | Generate = %7.2f tok/s\n",
warm_avg_prefill, warm_avg_generate);
}
fprintf(stderr,
"============================================================\n\n");
}

} // namespace
Expand Down
4 changes: 4 additions & 0 deletions ops/bench_matmul.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,15 +184,19 @@ void BenchAllMatMul() {
// QKV projection
BenchMatMul<BF16, BF16, BF16>(batch_size, 1152, 1536, kAdd, env);
BenchMatMul<BF16, SFP, BF16>(batch_size, 1152, 1536, kAdd, env);
BenchMatMul<BF16, Q4_0Stream, BF16>(batch_size, 1152, 1536, kAdd, env);
// FFN gate+up
BenchMatMul<BF16, BF16, BF16>(batch_size, 1152, 13824, kAdd, env);
BenchMatMul<BF16, SFP, BF16>(batch_size, 1152, 13824, kAdd, env);
BenchMatMul<BF16, Q4_0Stream, BF16>(batch_size, 1152, 13824, kAdd, env);
// FFN down
BenchMatMul<BF16, BF16, BF16>(batch_size, 6912, 1152, kAdd, env);
BenchMatMul<BF16, SFP, BF16>(batch_size, 6912, 1152, kAdd, env);
BenchMatMul<BF16, Q4_0Stream, BF16>(batch_size, 6912, 1152, kAdd, env);
// Logits / embedding
BenchMatMul<BF16, BF16, BF16>(batch_size, 1152, 262144, kAdd, env);
BenchMatMul<BF16, SFP, BF16>(batch_size, 1152, 262144, kAdd, env);
BenchMatMul<BF16, Q4_0Stream, BF16>(batch_size, 1152, 262144, kAdd, env);
}

PROFILER_PRINT_RESULTS();
Expand Down
Loading
Loading