Skip to content

compression: handle zero-length NUQ decompression #1062

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
Oct 19, 2026
Merged

compression: handle zero-length NUQ decompression #1062

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
2 changes: 2 additions & 0 deletions compression/nuq-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -756,6 +756,8 @@ class NuqCodec {
static HWY_INLINE void DecompressAndZeroPad(
D d, const PackedSpan<const NuqStream>& packed, size_t packed_ofs,
Raw* HWY_RESTRICT raw, size_t num) {
// num == 0 would make `num_groups - 1` below wrap to SIZE_MAX.
if (num == 0) return;
// If unaligned, load elements from the first group and update the args,
// from which we compute new tables/indices below.
size_t current_offset = packed_ofs;
Expand Down
27 changes: 27 additions & 0 deletions compression/nuq_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -507,6 +507,31 @@ struct TestEncDec {
void TestEncDecBF16() { hn::ForGEVectors<128, TestEncDec>()(BF16()); }
void TestEncDecF32() { hn::ForGEVectors<128, TestEncDec>()(float()); }

// Regression test: num == 0 must return immediately without any
// out-of-bounds access (previously `num_groups - 1` wrapped to SIZE_MAX).
struct TestZeroNum {
template <typename T, class D>
HWY_INLINE void operator()(T /*unused*/, D d) {
const hn::Repartition<float, D> df;
auto out = hwy::AllocateAligned<T>(kGroupSize);
auto nuq = hwy::AllocateAligned<NuqStream>(
hwy::RoundUpTo(NuqStream::PackedEnd(kGroupSize), hwy::VectorBytes()));
HWY_ASSERT(out && nuq);
// Sentinel-fill the output; the call must not touch it.
for (size_t i = 0; i < kGroupSize; ++i) {
out[i] = hwy::ConvertScalarTo<T>(-1.0f);
}
const auto nuq_span = MakeSpan(nuq.get(), kGroupSize);
NuqCodec::DecompressAndZeroPad(d, MakeConst(nuq_span), 0, out.get(), 0);
for (size_t i = 0; i < kGroupSize; ++i) {
HWY_ASSERT_EQ(hwy::ConvertScalarTo<T>(-1.0f), out[i]);
}
}
};

void TestZeroNumBF16() { hn::ForGEVectors<128, TestZeroNum>()(BF16()); }
void TestZeroNumF32() { hn::ForGEVectors<128, TestZeroNum>()(float()); }

// NOLINTNEXTLINE(google-readability-namespace-comments)
} // namespace HWY_NAMESPACE
} // namespace gcpp
Expand All @@ -529,6 +554,8 @@ HWY_EXPORT_AND_TEST_P(NuqTest, TestUnalignedOffsetF32);
HWY_EXPORT_AND_TEST_P(NuqTest, TestAllNibble);
HWY_EXPORT_AND_TEST_P(NuqTest, TestEncDecBF16);
HWY_EXPORT_AND_TEST_P(NuqTest, TestEncDecF32);
HWY_EXPORT_AND_TEST_P(NuqTest, TestZeroNumBF16);
HWY_EXPORT_AND_TEST_P(NuqTest, TestZeroNumF32);
#else
GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(NuqTest);
#endif // GEMMA_ENABLE_NUQ
Expand Down
Loading