Skip to content

E-2 on the AVX-2 path too: prefetch the i-quant expert rows (same switch, CPUs without AVX-512) #207

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
pipeob0 wants to merge 1 commit into from
Closed

E-2 on the AVX-2 path too: prefetch the i-quant expert rows (same switch, CPUs without AVX-512) #207

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
21 changes: 21 additions & 0 deletions src/kernels/cpu/iq_avx2.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include <immintrin.h>

#include <cmath>
#include <cstdlib>
#include <cstring>

namespace strata::kernels::cpu {
Expand Down Expand Up @@ -77,6 +78,23 @@ inline float hsum8(__m256 v) {
return _mm_cvtss_f32(s);
}

// E-2 (iq_avx512.cpp) on the AVX-2 path: the pool streams the expert rows from DRAM at ~25 GB/s (4 KB pages
// when large pages are refused), so ask for the bytes a few blocks before the decode needs them - 2048 B is
// two gate/up rows ahead, a row is ~1 KB. Same switch as the AVX-512 kernels: STRATA_IQ_PREFETCH is the
// distance in bytes, 0 = off, default 2048. Measured on a Zen 3 5700X3D (no AVX-512) on IQ3_S decode: -4% on
// the gate/up phase, +1.0 GB/s over the rows, -1.3% ms/round end to end. The non-temporal hint measured
// worse than T0 at the same distance, so this keeps T0.
const int prefetch_ahead = [] {
const char* v = std::getenv("STRATA_IQ_PREFETCH");
return v ? std::atoi(v) : 2048;
}();

inline void rows_ahead(const uint8_t* p) {
if (prefetch_ahead <= 0) return;
_mm_prefetch((const char*) p, _MM_HINT_T0);
_mm_prefetch((const char*) p + 64, _MM_HINT_T0);
}

// ---- per format: one 32-value half (values 64*j + 32*half .. +31) -> grid magnitudes, sign vector, scales
template <int TY> struct Fmt32;

Expand Down Expand Up @@ -165,6 +183,7 @@ inline void row_dot(const uint8_t* row, int nblocks, const block_q8_K* const* y,
for (int t = 0; t < NT; ++t) accf[t] = _mm256_setzero_ps();
for (int i = 0; i < nblocks; ++i) {
const uint8_t* blk = row + (size_t) i * Fmt32<TY>::bytes;
rows_ahead(blk + prefetch_ahead);
__m256i acci[NT];
for (int t = 0; t < NT; ++t) acci[t] = _mm256_setzero_si256();
for (int j = 0; j < 4; ++j) {
Expand Down Expand Up @@ -229,6 +248,7 @@ inline void row_dot_iq2xs(const uint8_t* row, int nblocks, const block_q8_K* con
for (int t = 0; t < NT; ++t) accf[t] = _mm256_setzero_ps();
for (int i = 0; i < nblocks; ++i) {
const uint8_t* blk = row + (size_t) i * 74;
rows_ahead(blk + prefetch_ahead);
// the 8 scale bytes -> 16 half-scales of 2*s+1, interleaved [a0, b0, a1, b1, ...] (ggml's unpack)
__m128i st = _mm_set1_epi64x((long long) u64(blk + 66));
st = _mm_unpacklo_epi8(_mm_and_si128(st, m4), _mm_and_si128(_mm_srli_epi16(st, 4), m4));
Expand Down Expand Up @@ -382,6 +402,7 @@ void iq4nl_rows(const uint8_t* w, size_t row_bytes, int n, const block_q8_0* con
for (int t = 0; t < NT; ++t) accf[t] = _mm256_setzero_ps();
for (int ib = 0; ib < nb; ++ib) {
const uint8_t* blk = row + (size_t) ib * sizeof(block_iq4_nl);
rows_ahead(blk + prefetch_ahead);
const __m128i bits = _mm_loadu_si128((const __m128i*) (blk + 2));
const __m128i lo = _mm_and_si128(bits, m4b); // values 0..15
const __m128i hi = _mm_and_si128(_mm_srli_epi16(bits, 4), m4b); // values 16..31
Expand Down