mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 06:15:58 -10:00
Exact two-pass QSA prompt selection (opt-in STRATA_SELECT_TWO_PASS=1)
Prompt-path block scores in single-pass TF32 with a per-block error bound, 3xTF32 rescoring only near each query's top-k cut-off, so qsa_block_topk returns the same cells in the same order, ties included. Off unless the variable is set; says once where it cannot run (HIP, Metal, SYCL, layer-major path). Adds qsa_two_pass_test, qsa_two_pass_status_test and the GPU parity test qsa_two_pass_parity. Ported from PR 4119 (bitaccesscomau); the log line of the per-prompt counters is shortened without dropping a field. l620 (sm_120, IQ3_S), greedy ids identical off vs on in every pair: prompt 128K +1.65% (n=12), 256K +4.10% (n=8), 32K -0.14% (noise). (cherry picked from commit bbf7c55f1351d913233250fe7869642f7c2b98a1)
This commit is contained in:
bitaccesscomau authored and Niko1221 committed
1 parent
f281e3fe26
commit
386f9959f1
11 files changed
+1860
-2
No files matched your search
@@ -403,6 +403,17 @@ if(STRATA_BUILD_TESTS)
|
||||
target_include_directories(share_rules_test PRIVATE include)
|
||||
add_test(NAME share_rules_test COMMAND share_rules_test)
|
||||
|
||||
# STRATA_SELECT_TWO_PASS: the bound between the single-pass and 3xTF32 QSA block scores and the cut-off rule that
|
||||
# keeps the selection (header-only, no GPU)
|
||||
add_executable(qsa_two_pass_test tests/prefill/qsa_two_pass_test.cpp)
|
||||
target_include_directories(qsa_two_pass_test PRIVATE include)
|
||||
add_test(NAME qsa_two_pass_test COMMAND qsa_two_pass_test)
|
||||
# STRATA_SELECT_TWO_PASS: the request's parse and the line it prints once where the scorer cannot run (no GPU)
|
||||
add_executable(qsa_two_pass_status_test tests/prefill/qsa_two_pass_status_test.cpp)
|
||||
target_include_directories(qsa_two_pass_status_test PRIVATE include)
|
||||
add_test(NAME qsa_two_pass_status_test COMMAND qsa_two_pass_status_test
|
||||
"${CMAKE_CURRENT_BINARY_DIR}/qsa_two_pass_status_test.log")
|
||||
|
||||
# STRATA_PREFILL_ROUTED_STREAM: the prompt stream's issuer walk and route gate, on a simulated link and consumer
|
||||
# (header-only, no GPU)
|
||||
add_executable(routed_stream_test tests/prefill/routed_stream_test.cpp)
|
||||
@@ -634,6 +645,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
|
||||
target_link_libraries(qsa_scores_parity PRIVATE strata_kernels CUDA::cudart)
|
||||
add_test(NAME qsa_scores_parity COMMAND qsa_scores_parity)
|
||||
set_tests_properties(qsa_scores_parity PROPERTIES SKIP_RETURN_CODE 77 ENVIRONMENT "STRATA_SCORES_THREAD=1")
|
||||
# STRATA_SELECT_TWO_PASS: the two-pass prompt scores' selection against the tensor-core scorer's, ids and the
|
||||
# scored-again blocks bitwise (sm_80+, synthetic, no model)
|
||||
add_executable(qsa_two_pass_parity src/kernels/qsa_two_pass_parity.cpp)
|
||||
target_include_directories(qsa_two_pass_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
|
||||
target_link_libraries(qsa_two_pass_parity PRIVATE strata_kernels CUDA::cudart)
|
||||
add_test(NAME qsa_two_pass_parity COMMAND qsa_two_pass_parity)
|
||||
set_tests_properties(qsa_two_pass_parity PROPERTIES SKIP_RETURN_CODE 77 ENVIRONMENT "STRATA_SELECT_TWO_PASS=1")
|
||||
endif()
|
||||
# QSAROWS (docs/QSAROWS.md): what the T rows of a verification window cost in a QSA layer's decode kernels at a long
|
||||
# context (block scores, top-k, attention over each row's selection, indexer append), row by row; GPU, synthetic, no
|
||||
|
||||
@@ -1119,6 +1119,18 @@ measured. "Exact" means the 10-prompt and 4K/20K identity gate with the switch o
|
||||
`STRATA_DEQUANT_WIDE` (wide dequantization of Q4_K, Q5_K, Q6_K and IQ4_XS for the prompt's dense weights) is on by default on
|
||||
sm_86 and sm_61, bit-exact. `STRATA_DEQUANT_WIDE=0` turns it off.
|
||||
|
||||
`STRATA_SELECT_TWO_PASS=1` (opt-in; sm_80+, the prompt path's tensor-core block scorer): every block is scored once in
|
||||
single-pass TF32 with a per-block error bound, and again in the usual 3xTF32 only where the bound reaches the query's
|
||||
top-k cut-off (a query the bound cannot settle has its 16-query tile scored again in full), so the selected cells, their
|
||||
order and the ties are exactly the 3xTF32 scorer's (the bound and the rule: `include/strata/kernels/qsa_two_pass.hpp`).
|
||||
On an RTX 3090 (sm_86, PCIe 3.0, 250 W; the v0.2.6 release defaults, 2 pairs of 2 runs) a 256K prompt read 5.2% faster
|
||||
(1,887 to 1,984 tokens/s), a 128K prompt 2.2% faster and a 32K prompt the same within 1%; decode and the answers did not
|
||||
change. Calls whose last query has `STRATA_SELECT_TWO_PASS_MIN` blocks or more take it (default 4096, 16K cells).
|
||||
`STRATA_SELECT_TWO_PASS_CHECK=1` (a slow diagnostic) runs the 3xTF32 scorer beside it and prints, after each prompt, the
|
||||
queries whose ids differ (must be 0). Where the scorer cannot run (HIP, Metal and SYCL builds, a card below sm_80 or
|
||||
`STRATA_QSA_WARP`, the layer-major prompt path) a set switch prints one line saying it is not used there, and the
|
||||
selection is unchanged.
|
||||
|
||||
`STRATA_PREFILL_ROUTED_STREAM=1` (opt-in; the streamed prompt walk of a chunk of 1,024 tokens or more, one GPU, the
|
||||
issuer thread): the issuer keeps at most `STRATA_PREFILL_ROUTED_STREAM_DEPTH` copies (default 16) queued ahead of the
|
||||
PCIe link, each layer's routing is handed to it as soon as the layer's router has run, and an expert not yet copied that
|
||||
|
||||
@@ -94,4 +94,33 @@ void qsa_block_scores_long(const float* pooled, const float* dead, const float*
|
||||
/// path: pass null).
|
||||
int64_t qsa_topk_multi_scratch_ints(int64_t nq);
|
||||
|
||||
/// STRATA_SELECT_TWO_PASS=1 (the prompt path, CUDA sm_80+; the bound and the rule: qsa_two_pass.hpp). A scores row
|
||||
/// that qsa_block_topk turns into exactly the ids of qsa_block_scores_tc's row: every block is scored once in
|
||||
/// single-pass TF32, the blocks whose error bound reaches the query's top-k cut-off are scored again with
|
||||
/// qsa_block_scores_tc's own 3xTF32 arithmetic (bit for bit), and a query whose bound cannot be settled has its
|
||||
/// 16-query tile scored again in full. The kept blocks hold their single-pass score, so the row is NOT
|
||||
/// qsa_block_scores_tc's: it is for qsa_block_topk only. `scratch` (qsa_two_pass_scratch_bytes(nq_cap, max_blocks),
|
||||
/// 16-byte aligned) holds the key norms qsa_two_pass_keys wrote for this layer and chunk. False (nothing launched):
|
||||
/// the switch is off, HIP, a card or STRATA_QSA_WARP where qsa_block_scores_tc would not run its tensor-core kernel,
|
||||
/// another geometry, nq > nq_cap, too little scratch.
|
||||
bool qsa_two_pass_on(); ///< STRATA_SELECT_TWO_PASS=1, read once (HIP: false; prints once where unused)
|
||||
bool qsa_two_pass_check_on(); ///< STRATA_SELECT_TWO_PASS_CHECK=1, read once (a diagnostic, see qsa_two_pass_check)
|
||||
int64_t qsa_two_pass_min_blocks(); ///< STRATA_SELECT_TWO_PASS_MIN (blocks), read once; default two_pass::kMinBlocksDefault
|
||||
uint64_t qsa_two_pass_scratch_bytes(int64_t nq_cap, int64_t max_blocks);
|
||||
/// Once per layer and chunk, after the indexer appends and before the calls: the norms of the pooled keys [0, n_blocks).
|
||||
bool qsa_two_pass_keys(const float* pooled, int64_t n_blocks, int64_t nq_cap, int64_t max_blocks, void* scratch,
|
||||
uint64_t scratch_bytes, void* stream);
|
||||
bool qsa_block_scores_two_pass(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps,
|
||||
int64_t nq, int64_t nq_cap, int64_t max_blocks, const QsaShapes& s, float* scores,
|
||||
void* scratch, uint64_t scratch_bytes, void* stream, int64_t active_blocks);
|
||||
/// STRATA_SELECT_TWO_PASS_CHECK: right after a two-pass call and its qsa_block_topk, against today's row and ids for
|
||||
/// the same queries (ref_scores from qsa_block_scores_tc, ref_ids from qsa_block_topk on them): counts the queries
|
||||
/// whose ids differ, the scored-again blocks whose score is not today's bit for bit, and the kept blocks whose
|
||||
/// single-pass score lies outside the bound (two_pass::Stat).
|
||||
void qsa_two_pass_check(const float* scores, const int32_t* ids, const float* ref_scores, const int32_t* ref_ids,
|
||||
const int32_t* steps, int64_t nq, int64_t nq_cap, int64_t max_blocks, int64_t cap,
|
||||
void* scratch, void* stream);
|
||||
/// The counters (two_pass::kStatCount of them, this device), read on `stream` (synchronized); reset: zeroed after.
|
||||
bool qsa_two_pass_stats(unsigned long long* out, bool reset, void* stream);
|
||||
|
||||
} // namespace strata::kernels
|
||||
@@ -0,0 +1,141 @@
|
||||
// include/strata/kernels/qsa_two_pass.hpp - STRATA_SELECT_TWO_PASS: the prompt path's QSA block scores in two
|
||||
// precisions with the same selection. Header-only and host/device: the CUDA scorer (qsa_select.cu) and the GPU-free
|
||||
// test (tests/prefill/qsa_two_pass_test.cpp) read the same constants.
|
||||
//
|
||||
// Today (qsa_block_scores_tc, sm_80+) every block score is 3xTF32: per indexer head h and pooled key k, with
|
||||
// qh = tf32(q), ql = tf32(q - qh), kh = tf32(k), kl = tf32(k - kh) (cvt.rna), the chain of 16 m16n8k8 steps adds
|
||||
// ql*kh, qh*kl, qh*kh into one FP32 accumulator (48 MMAs), and the score is relu(head 0) + ... + relu(head 3).
|
||||
// The two-pass scorer computes S1 (qh*kh only, 16 MMAs per head) for every block, bounds |S1 - S3| per block, and
|
||||
// computes S3 again, bit for bit as today, only for the blocks whose bound reaches the top-k cut-off. Every other
|
||||
// block keeps S1, which lies strictly on the same side of the cut-off as its S3, so qsa_block_topk returns exactly
|
||||
// today's cells in today's order (the threshold key, the count above it and the tied cells are unchanged).
|
||||
//
|
||||
// THE BOUND (per query q, block b; derived for the PR, conservative on purpose):
|
||||
// S1_h - S3_h = -(qh.kl + ql.kh) + (accumulation error of S1) - (accumulation error of S3)
|
||||
// |qh.kl + ql.kh| <= |qh| |kl| + |ql| |kh| (Cauchy-Schwarz, exact TF32 operands)
|
||||
// accumulation: each MMA's FP32 result is taken to be within kUlps * 2^-23 of the absolute sum of its 9 addends
|
||||
// (the accumulator and 8 exact TF32 products). A truncating alignment with no guard bits gives at most 9 such
|
||||
// ulps (8 truncated addends plus the final rounding), so 16 is a margin, not a fit. Over a chain of n MMAs the
|
||||
// error is at most n * kUlps * 2^-23 * A, A = the chain's absolute sum of products:
|
||||
// A(S1) <= |qh| |kh|, A(S3) <= |qh| |kh| + |qh| |kl| + |ql| |kh|
|
||||
// relu is 1-Lipschitz; the head sums ((r0 + r1) + r2) + r3 add at most 3 * 2^-24 relative each (folded into kG3).
|
||||
// Summed over the four heads, with a = sum_h |qh_h|, c = sum_h |ql_h| (per query) and kh = |kh_b|, kl = |kl_b|:
|
||||
// eps_b = kl * P + kh * Q + kAbs, P = (1 + kG3) a s, Q = ((1 + kG3) c + (kG1 + kG3) a) s, s = kSlack
|
||||
// kAbs (2^-100) covers flushed subnormals (inputs, products, accumulators); kSlack the second-order terms.
|
||||
// All norms are upper bounds (double, rounded up), eps is rounded up and the interval [S1 - eps, S1 + eps] is
|
||||
// rounded outwards (CUDA __f*_rd / __f*_ru), so the interval is an enclosure of S3 under the stated MMA model.
|
||||
//
|
||||
// THE CUT-OFF RULE (per query): T = the width-th largest S3 cell key, the threshold qsa_block_topk computes. With
|
||||
// T1 = the width-th largest S1 and E = the largest eps over the candidate blocks, T lies in [T1 - E, T1 + E]: the
|
||||
// width cells with S1 >= T1 have S3 >= T1 - E, and every block with S3 > T1 + E has S1 > T1. A block with
|
||||
// S1 - eps > T1 + E is above T for certain, one with S1 + eps < T1 - E below it; every other block (including every
|
||||
// block whose S3 equals T, so the ties) is scored again with today's 3xTF32 arithmetic.
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
#if defined(__CUDACC__) || defined(__HIPCC__)
|
||||
#define STRATA_TP_HD __host__ __device__
|
||||
#else
|
||||
#define STRATA_TP_HD
|
||||
#endif
|
||||
|
||||
namespace strata::kernels::two_pass {
|
||||
|
||||
constexpr int kSteps = 16; ///< K = 128 in m16n8k8 steps, one chain per head
|
||||
constexpr double kUlps = 16.0; ///< accumulation allowance per MMA, FP32 ulps of its absolute sum
|
||||
constexpr double kG1 = kSteps * kUlps * 0x1p-23; ///< the single-pass chain: 16 MMAs
|
||||
constexpr double kG3 = 3 * kSteps * kUlps * 0x1p-23 + 0x1p-21; ///< the 3xTF32 chain: 48 MMAs, plus the head sums
|
||||
constexpr double kSlack = 1.0 + 0x1p-10;
|
||||
constexpr float kAbs = 0x1p-100f;
|
||||
constexpr float kMaxEps = 1e30f; ///< larger bounds (or non-finite norms) mean: score this query again
|
||||
|
||||
/// Candidate blocks per query held by the band kernel (shared memory); more, and the query is scored again in full.
|
||||
constexpr int kCandCap = 5120;
|
||||
/// Blocks re-scored per query and per column residue (b % 8): eight lists of this many; more is a full re-score.
|
||||
constexpr int kResCap = 128;
|
||||
constexpr int kSampleStride = 32; ///< the candidate threshold comes from every 32nd block, as the
|
||||
///< prompt top-k's sample kernel does
|
||||
/// The calls (a batch of up to 256 queries) that take the two-pass path by default: the batch's last query has at
|
||||
/// least this many blocks (16,384 cells). STRATA_SELECT_TWO_PASS_MIN=<blocks> moves it (0: every call).
|
||||
constexpr int64_t kMinBlocksDefault = 4096;
|
||||
|
||||
/// The per-query coefficients of eps_b = kl_b * P + kh_b * Q (+ kAbs), from the head norm sums a and c.
|
||||
STRATA_TP_HD inline void coefficients(double a, double c, double& P, double& Q) {
|
||||
P = (1.0 + kG3) * a * kSlack;
|
||||
Q = ((1.0 + kG3) * c + (kG1 + kG3) * a) * kSlack;
|
||||
}
|
||||
|
||||
/// qsa_select.cu's order_key: the float's order as an unsigned key (NaN lowest, -0 as +0).
|
||||
STRATA_TP_HD inline uint32_t order_key_host(float s) {
|
||||
const float v = s + 0.0f;
|
||||
if (!(v == v)) return 0u;
|
||||
uint32_t b;
|
||||
std::memcpy(&b, &v, 4);
|
||||
return (b & 0x80000000u) ? ~b : (b | 0x80000000u);
|
||||
}
|
||||
/// The float of a key order_key produced (the inverse, for keys of non-NaN values).
|
||||
STRATA_TP_HD inline float key_float_host(uint32_t k) {
|
||||
const uint32_t b = (k & 0x80000000u) ? (k & 0x7fffffffu) : ~k;
|
||||
float f;
|
||||
std::memcpy(&f, &b, 4);
|
||||
return f;
|
||||
}
|
||||
|
||||
/// cvt.rna.tf32.f32 on the host: round to nearest, ties away from zero, 10 explicit mantissa bits (low 13 bits 0).
|
||||
inline float tf32_rna_host(float x) {
|
||||
uint32_t b;
|
||||
std::memcpy(&b, &x, 4);
|
||||
if ((b & 0x7f800000u) == 0x7f800000u) return x; // Inf and NaN unchanged
|
||||
b = (b + 0x1000u) & 0xffffe000u; // magnitude rounding; a carry moves into the exponent
|
||||
float r;
|
||||
std::memcpy(&r, &b, 4);
|
||||
return r;
|
||||
}
|
||||
|
||||
/// Scratch the two-pass scorer takes from the caller (bytes, 256-aligned parts): the key norms of every block
|
||||
/// (one layer, one chunk) and, per call of at most nq_cap queries, the coefficients, modes, residue counts, re-score
|
||||
/// lists and the 16-query tile flags.
|
||||
struct Layout {
|
||||
uint64_t kn = 0, qc = 0, mode = 0, cnt = 0, ids = 0, tiles = 0, total = 0;
|
||||
};
|
||||
inline Layout layout(int64_t nq_cap, int64_t max_blocks) {
|
||||
auto up = [](uint64_t v) { return (v + 255u) & ~uint64_t(255u); };
|
||||
Layout l;
|
||||
uint64_t o = 0;
|
||||
l.kn = o; o += up((uint64_t) max_blocks * 8u);
|
||||
l.qc = o; o += up((uint64_t) nq_cap * 8u);
|
||||
l.mode = o; o += up((uint64_t) nq_cap * 4u);
|
||||
l.cnt = o; o += up((uint64_t) nq_cap * 8u * 4u);
|
||||
l.ids = o; o += up((uint64_t) nq_cap * 8u * (uint64_t) kResCap * 4u);
|
||||
l.tiles = o; o += up((uint64_t) ((nq_cap + 15) / 16) * 4u);
|
||||
l.total = o;
|
||||
return l;
|
||||
}
|
||||
|
||||
/// The engine's counters (qsa_two_pass_stats): what the scorer did, and with STRATA_SELECT_TWO_PASS_CHECK=1 what
|
||||
/// the comparison against today's scorer found.
|
||||
enum Stat : int {
|
||||
kStCalls = 0, ///< two-pass calls
|
||||
kStQueries, ///< their queries with a sparse selection (n_kv > width)
|
||||
kStIdentity, ///< their queries whose selection is every cell (nothing to score)
|
||||
kStFull, ///< queries flagged for a full re-score (no bound, a list full, too few cells above F)
|
||||
kStFullTiles, ///< 16-query tiles scored again in full (3xTF32, today's kernel; every query of the tile)
|
||||
kStBand, ///< blocks scored again (3xTF32) next to the cut-off
|
||||
kStBandMax, ///< the most blocks scored again for one query
|
||||
kStCand, ///< candidate blocks (the band kernel's shared list)
|
||||
kStCandMax,
|
||||
kStBlocks, ///< blocks scored once (single-pass) in two-pass calls
|
||||
kChkQueries, ///< check: queries compared (ids)
|
||||
kChkIdQueries, ///< check: queries whose ids differ from today's
|
||||
kChkIdEntries, ///< check: id entries that differ
|
||||
kChkRescored, ///< check: scored-again blocks compared bitwise with today's score
|
||||
kChkRescoredDiff, ///< check: ... that differ (must be 0)
|
||||
kChkKept, ///< check: blocks that kept S1 (compared against eps)
|
||||
kChkKeptOut, ///< check: ... with |S1 - S3| > eps (must be 0)
|
||||
kChkRatioMax, ///< check: the largest |S1 - S3| / eps seen, as float bits
|
||||
kStatCount
|
||||
};
|
||||
|
||||
} // namespace strata::kernels::two_pass
|
||||
@@ -0,0 +1,41 @@
|
||||
#pragma once
|
||||
// STRATA_SELECT_TWO_PASS: the request, read once, and the line a set switch prints where the two-pass scorer cannot
|
||||
// run. The scorer runs only in the CUDA prompt path's chunked walk, on a card where qsa_block_scores_tc runs its
|
||||
// tensor-core kernel (sm_80+, STRATA_QSA_WARP not set). Elsewhere (HIP, Metal, SYCL, an older card or the warp
|
||||
// scorer, the layer-major prompt path) the selection runs as without the switch; that was silent, so the switch read
|
||||
// as active. Header-only, so the rule and the line are tested without a GPU (qsa_two_pass_status_test).
|
||||
#include <atomic>
|
||||
#include <cstdio>
|
||||
|
||||
namespace strata::kernels::two_pass {
|
||||
|
||||
// The switch's own parse: exactly "1" is on (as the two-pass scorer read it before).
|
||||
inline bool requested_from(const char* v) { return v != nullptr && v[0] == '1' && v[1] == '\0'; }
|
||||
|
||||
class RequestStatus {
|
||||
public:
|
||||
explicit RequestStatus(const char* flag) : requested_(requested_from(flag)) {}
|
||||
|
||||
bool requested() const { return requested_; }
|
||||
|
||||
// `not_here` null: this path can run the scorer, so the result is the request (whether a call takes it still
|
||||
// depends on the call: its size, geometry and scratch). Otherwise the path cannot run it: false, and the first
|
||||
// such query prints one line naming the path. A supported query never uses up that line.
|
||||
bool on(const char* not_here, FILE* log = stderr) {
|
||||
if (!requested_) return false;
|
||||
if (not_here == nullptr) return true;
|
||||
if (!warned_.exchange(true)) {
|
||||
std::fprintf(log, "strata: STRATA_SELECT_TWO_PASS=1 is not used on %s: the two-pass scorer runs only in "
|
||||
"the CUDA prompt path's tensor-core selection (sm_80+); the selection is unchanged\n",
|
||||
not_here);
|
||||
std::fflush(log);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
private:
|
||||
const bool requested_;
|
||||
std::atomic<bool> warned_{false};
|
||||
};
|
||||
|
||||
} // namespace strata::kernels::two_pass
|
||||
@@ -5,6 +5,8 @@
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include "strata/kernels/qsa_select.hpp"
|
||||
#include "strata/kernels/qsa_two_pass.hpp"
|
||||
#include "strata/kernels/qsa_two_pass_status.hpp"
|
||||
#include "strata/kernels/gfx_arch.hpp"
|
||||
#include "strata/core/amd_community.hpp"
|
||||
|
||||
@@ -3143,4 +3145,764 @@ void qsa_block_scores_long(const float* pooled, const float* dead, const float*
|
||||
qsa_block_scores(pooled, dead, q_idx, steps, nq, max_blocks, s, scores, stream, active_blocks);
|
||||
}
|
||||
|
||||
// ==== STRATA_SELECT_TWO_PASS (the prompt path's block scores in two precisions; qsa_two_pass.hpp has the bound) ====
|
||||
// qsa_block_scores_tc computes every block score in 3xTF32. Here every block is scored once in single-pass TF32
|
||||
// (S1), a per-block bound eps encloses today's 3xTF32 score S3 in [S1 - eps, S1 + eps], and only the blocks whose
|
||||
// interval reaches the query's top-k cut-off are scored again, with block_scores_tc_kernel's arithmetic, bit for bit.
|
||||
// qsa_block_topk then reads a row whose every block lies on the same side of the threshold as in today's row and
|
||||
// whose blocks at the threshold hold today's exact scores: the same threshold, the same ties, the same cells. The row
|
||||
// itself is not today's (kept blocks hold S1), so it is for qsa_block_topk only: nothing else reads m.sel_scores.
|
||||
#if !defined(__HIPCC__)
|
||||
__device__ unsigned long long g_tp_stats[two_pass::kStatCount];
|
||||
|
||||
namespace {
|
||||
namespace tp = strata::kernels::two_pass;
|
||||
|
||||
constexpr int TP1_ITER = 8; // 32-block tiles per single-pass CTA (its query fragments are loaded once)
|
||||
constexpr int TPB_T = 512; // the band kernel's threads (one query per CTA)
|
||||
static_assert(TPB_T == IDX_HEADS * IDX_DIM, "the band kernel reads one query element per thread");
|
||||
|
||||
__device__ __forceinline__ float tp_inf() { return __int_as_float(0x7f800000); }
|
||||
__device__ __forceinline__ float tp_key_float(uint32_t k) { // order_key's inverse (keys of non-NaN floats)
|
||||
return __uint_as_float((k & 0x80000000u) ? (k & 0x7fffffffu) : ~k);
|
||||
}
|
||||
|
||||
// The single-pass enclosure of block b's 3xTF32 score from its single-pass score s: [lo, hi], eps rounded up and the
|
||||
// ends rounded outwards. False: no bound (a non-finite score, norm or coefficient, or a bound past kMaxEps).
|
||||
__device__ __forceinline__ bool tp_interval(float s, float2 kn, float P, float Q, float& lo, float& hi, float& e) {
|
||||
e = __fmaf_ru(P, kn.y, __fmaf_ru(Q, kn.x, tp::kAbs));
|
||||
if (!(e <= tp::kMaxEps) || !(fabsf(s) <= FLT_MAX)) return false;
|
||||
lo = __fsub_rd(s, e);
|
||||
hi = __fadd_ru(s, e);
|
||||
return true;
|
||||
}
|
||||
|
||||
// The key norms (upper bounds): per pooled block |tf32(k)| and |tf32(k - tf32(k))|, the split block_scores_tc_kernel
|
||||
// makes, in double with every operation rounded up; a non-finite norm becomes +inf (its blocks are never bounded).
|
||||
__global__ void __launch_bounds__(256) tp_keys_kernel(const float* __restrict__ pooled, int64_t n_blocks,
|
||||
float2* __restrict__ kn) {
|
||||
const int64_t b = (int64_t) blockIdx.x * 8 + (threadIdx.x >> 5);
|
||||
if (b >= n_blocks) return;
|
||||
const int lane = threadIdx.x & 31;
|
||||
const float4 v = reinterpret_cast<const float4*>(pooled + b * IDX_DIM)[lane];
|
||||
const float x[4] = {v.x, v.y, v.z, v.w};
|
||||
double sh = 0.0, sl = 0.0;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
const float h = __uint_as_float(tf32_hi(x[j]));
|
||||
const float l = __uint_as_float(tf32_hi(x[j] - h));
|
||||
sh = __dadd_ru(sh, __dmul_ru((double) h, (double) h));
|
||||
sl = __dadd_ru(sl, __dmul_ru((double) l, (double) l));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
sh = __dadd_ru(sh, __shfl_xor_sync(0xffffffffu, sh, o));
|
||||
sl = __dadd_ru(sl, __shfl_xor_sync(0xffffffffu, sl, o));
|
||||
}
|
||||
if (lane == 0) {
|
||||
float kh = __double2float_ru(__dsqrt_ru(sh)), kl = __double2float_ru(__dsqrt_ru(sl));
|
||||
if (!(kh <= FLT_MAX)) kh = tp_inf();
|
||||
if (!(kl <= FLT_MAX)) kl = tp_inf();
|
||||
kn[b] = make_float2(kh, kl);
|
||||
}
|
||||
}
|
||||
|
||||
// Pass 1: S1 = sum over the heads of relu(tf32(q) . tf32(k)), one MMA per step (block_scores_tc_kernel makes three).
|
||||
// Warp h is head h and keeps its 16 steps of A fragments in registers; a CTA is one 16-query tile (blockIdx.x, so the
|
||||
// CTAs that run together read the same keys) and TP1_ITER tiles of 32 blocks. Blocks < n_bid only, as today; the
|
||||
// tail block is block_scores_tail_kernel's. Also clears the tile flags the band kernel sets.
|
||||
__global__ void __launch_bounds__(128) tp_score1_kernel(const float* __restrict__ pooled,
|
||||
const float* __restrict__ q_idx,
|
||||
const int32_t* __restrict__ steps, int64_t nq,
|
||||
int64_t max_blocks, int64_t reach, float* __restrict__ out,
|
||||
int* __restrict__ tiles) {
|
||||
__shared__ __align__(16) uint32_t sK[TC_NB * TC_KS];
|
||||
__shared__ float red[IDX_HEADS][TC_QT][TC_NB + 1];
|
||||
__shared__ int s_nbid[TC_QT];
|
||||
const int t = threadIdx.x, lane = t & 31, h = t >> 5, gid = lane >> 2, tig = lane & 3;
|
||||
const int64_t q0 = (int64_t) blockIdx.x * TC_QT;
|
||||
if (blockIdx.y == 0 && t == 0) tiles[blockIdx.x] = 0;
|
||||
if (t < TC_QT) s_nbid[t] = q0 + t < nq ? steps[(q0 + t) * kStepCount + kStepNBid] : 0;
|
||||
uint32_t a[16][4]; // rows gid and gid + 8, k = 8s + tig and 8s + tig + 4 (the m16n8k8 A fragment)
|
||||
{
|
||||
const bool r0 = q0 + gid < nq, r1 = q0 + gid + 8 < nq;
|
||||
const float* qa = q_idx + (r0 ? q0 + gid : 0) * (IDX_HEADS * IDX_DIM) + h * IDX_DIM + tig;
|
||||
const float* qb = q_idx + (r1 ? q0 + gid + 8 : 0) * (IDX_HEADS * IDX_DIM) + h * IDX_DIM + tig;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 16; ++s) {
|
||||
a[s][0] = tf32_hi(r0 ? qa[8 * s] : 0.f);
|
||||
a[s][1] = tf32_hi(r1 ? qb[8 * s] : 0.f);
|
||||
a[s][2] = tf32_hi(r0 ? qa[8 * s + 4] : 0.f);
|
||||
a[s][3] = tf32_hi(r1 ? qb[8 * s + 4] : 0.f);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
int hi_nbid = 0;
|
||||
for (int i = 0; i < TC_QT; ++i) hi_nbid = max(hi_nbid, s_nbid[i]);
|
||||
for (int it = 0; it < TP1_ITER; ++it) {
|
||||
const int64_t b0 = ((int64_t) blockIdx.y * TP1_ITER + it) * TC_NB;
|
||||
if (b0 >= reach || b0 >= hi_nbid) break;
|
||||
__syncthreads(); // the previous tile's sK and red reads are done
|
||||
for (int i = t; i < TC_NB * IDX_DIM / 4; i += 128) {
|
||||
const int r = i / (IDX_DIM / 4), c = i % (IDX_DIM / 4);
|
||||
float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
if (b0 + r < hi_nbid) v = reinterpret_cast<const float4*>(pooled + (b0 + r) * IDX_DIM)[c];
|
||||
*reinterpret_cast<uint4*>(sK + r * TC_KS + c * 4) =
|
||||
make_uint4(tf32_hi(v.x), tf32_hi(v.y), tf32_hi(v.z), tf32_hi(v.w));
|
||||
}
|
||||
__syncthreads();
|
||||
float acc[4][4];
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 4; ++n) acc[n][0] = acc[n][1] = acc[n][2] = acc[n][3] = 0.f;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 16; ++s) {
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 4; ++n) {
|
||||
const uint32_t* kr = sK + (n * 8 + gid) * TC_KS + 8 * s + tig;
|
||||
const uint32_t b[2] = {kr[0], kr[4]};
|
||||
mma_tf32(acc[n], a[s], b);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 4; ++n) {
|
||||
const int c = n * 8 + 2 * tig;
|
||||
red[h][gid][c] = acc[n][0] > 0.0f ? acc[n][0] : 0.0f;
|
||||
red[h][gid][c + 1] = acc[n][1] > 0.0f ? acc[n][1] : 0.0f;
|
||||
red[h][gid + 8][c] = acc[n][2] > 0.0f ? acc[n][2] : 0.0f;
|
||||
red[h][gid + 8][c + 1] = acc[n][3] > 0.0f ? acc[n][3] : 0.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
for (int e = t; e < TC_QT * TC_NB; e += 128) {
|
||||
const int r = e / TC_NB, c = e % TC_NB;
|
||||
const int64_t qi = q0 + r, b = b0 + c;
|
||||
if (qi >= nq || b >= s_nbid[r] || b >= max_blocks) continue;
|
||||
float score = 0.0f;
|
||||
#pragma unroll
|
||||
for (int h2 = 0; h2 < IDX_HEADS; ++h2) score += red[h2][r][c];
|
||||
out[qi * max_blocks + b] = score;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// block_topk_kernel's threshold over a list in shared memory: the largest key k such that the weighted entries with
|
||||
// a key >= k reach `width`, 8 bits a pass with that kernel's digit rule (the largest digit whose cells at or above it
|
||||
// reach the width, else 0). unit: every entry weighs 1 (the sample); else a listed block weighs R and the tail block
|
||||
// n_bid w_tail. Every thread of the CTA calls it; it ends synchronized.
|
||||
__device__ uint32_t tp_select(const uint32_t* __restrict__ key, const int32_t* __restrict__ blk, int n, int width,
|
||||
bool unit, int64_t n_bid, int w_tail, int* hist, int* s_misc) {
|
||||
const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
|
||||
uint32_t prefix = 0;
|
||||
int above = 0;
|
||||
for (int shift = 24; shift >= 0; shift -= 8) {
|
||||
for (int i = t; i < 256; i += TPB_T) hist[i] = 0;
|
||||
__syncthreads();
|
||||
const uint32_t hi_mask = shift == 24 ? 0u : (0xffffffffu << (shift + 8));
|
||||
for (int i0 = 0; i0 < n; i0 += TPB_T) { // the same trip count in every thread: whole warps match
|
||||
const int i = i0 + t;
|
||||
const uint32_t k = i < n ? key[i] : 0u;
|
||||
const bool hit = i < n && (k & hi_mask) == (prefix & hi_mask);
|
||||
const bool tail = hit && !unit && blk[i] == n_bid;
|
||||
const int d = (int) ((k >> shift) & 255u);
|
||||
const int g = hit && !tail ? d : 256;
|
||||
const unsigned grp = __match_any_sync(0xffffffffu, g); // one add per digit and warp
|
||||
if (g < 256 && lane == __ffs(grp) - 1) atomicAdd(&hist[d], (unit ? 1 : R) * __popc(grp));
|
||||
if (tail) atomicAdd(&hist[d], w_tail);
|
||||
}
|
||||
__syncthreads();
|
||||
if (warp == 0) { // the digit: lane l holds digits 8l .. 8l + 7
|
||||
int v[8], local = 0;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j) { v[j] = hist[lane * 8 + j]; local += v[j]; }
|
||||
int suf = local; // cells in this lane's digits and every higher one
|
||||
#pragma unroll
|
||||
for (int o = 1; o < 32; o <<= 1) {
|
||||
const int y = __shfl_down_sync(0xffffffffu, suf, o);
|
||||
if (lane + o < 32) suf += y;
|
||||
}
|
||||
int cum = above + suf - local, best = -1, best_cum = 0;
|
||||
#pragma unroll
|
||||
for (int j = 7; j >= 0; --j) {
|
||||
const int dd = lane * 8 + j;
|
||||
if (best < 0 && dd >= 1 && cum + v[j] >= width) { best = dd; best_cum = cum; }
|
||||
cum += v[j];
|
||||
}
|
||||
const unsigned m = __ballot_sync(0xffffffffu, best >= 0);
|
||||
const int top = m ? 31 - __clz((int) m) : 0;
|
||||
const int dsel = __shfl_sync(0xffffffffu, best, top), csel = __shfl_sync(0xffffffffu, best_cum, top);
|
||||
const int all = __shfl_sync(0xffffffffu, suf, 0), h0 = __shfl_sync(0xffffffffu, v[0], 0);
|
||||
if (lane == 0) {
|
||||
s_misc[0] = m ? dsel : 0;
|
||||
s_misc[1] = m ? csel : above + all - h0; // digit 0: every cell of digits 1..255 is above it
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
prefix |= (uint32_t) s_misc[0] << shift;
|
||||
above = s_misc[1];
|
||||
__syncthreads();
|
||||
}
|
||||
return prefix;
|
||||
}
|
||||
|
||||
// Per query: the blocks to score again. 1. the coefficients P, Q (qsa_two_pass.hpp) from the query's head norms;
|
||||
// 2. a candidate threshold F from every 32nd block (only when the blocks outnumber the list): the ks-th largest
|
||||
// sampled S1 - eps; 3. one read of every block: the candidates (S1 + eps >= F), the cells with S1 - eps >= F (at least
|
||||
// `width` of them, so the threshold T of today's row and the single-pass one T1 are both >= F and every block that
|
||||
// matters below is a candidate) and E = the largest candidate eps; 4. T1 over the candidates; 5. T lies in
|
||||
// [max(T1 - E, F), T1 + E]: a candidate whose interval reaches it (and is not the exact tail block) is listed for a
|
||||
// re-score, by its column residue b % 8. A query that cannot be settled (no bound for some candidate, a list or a
|
||||
// residue list full, too few cells above F) is flagged: its 16-query tile is scored again in full.
|
||||
__global__ void __launch_bounds__(TPB_T) tp_band_kernel(const float* __restrict__ scores,
|
||||
const float* __restrict__ q_idx,
|
||||
const int32_t* __restrict__ steps, int64_t max_blocks,
|
||||
const float2* __restrict__ kn, float2* __restrict__ qc,
|
||||
int* __restrict__ mode, int* __restrict__ cnt,
|
||||
int* __restrict__ ids, int* __restrict__ tiles) {
|
||||
extern __shared__ __align__(16) uint32_t tp_sm[];
|
||||
uint32_t* c_key = tp_sm; // [kCandCap] S1 keys (first: the sample)
|
||||
int32_t* c_blk = reinterpret_cast<int32_t*>(tp_sm + tp::kCandCap); // [kCandCap] their blocks
|
||||
__shared__ int hist[256];
|
||||
__shared__ int s_misc[2], s_n, s_bad, s_wf[TPB_T / 32], s_cnt[8];
|
||||
__shared__ float s_e[TPB_T / 32], s_P, s_Q, s_E;
|
||||
__shared__ double s_hn[TPB_T / 32][2];
|
||||
const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
|
||||
const int64_t qi = blockIdx.x;
|
||||
const int32_t* st = steps + qi * kStepCount;
|
||||
const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth];
|
||||
if (qi == 0 && t == 0) atomicAdd(&g_tp_stats[tp::kStCalls], 1ull);
|
||||
if (n_kv <= width) { // every cell is selected: qsa_block_topk reads no score
|
||||
if (t == 0) { mode[qi] = 0; atomicAdd(&g_tp_stats[tp::kStIdentity], 1ull); }
|
||||
return;
|
||||
}
|
||||
{ // 1. a = sum_h |tf32(q_h)|, c = sum_h |tf32(q_h - tf32(q_h))|, rounded up; warps 4h .. 4h + 3 hold head h
|
||||
const float x = q_idx[qi * IDX_HEADS * IDX_DIM + t];
|
||||
const float hx = __uint_as_float(tf32_hi(x));
|
||||
const float lx = __uint_as_float(tf32_hi(x - hx));
|
||||
double sh = __dmul_ru((double) hx, (double) hx), sl = __dmul_ru((double) lx, (double) lx);
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
sh = __dadd_ru(sh, __shfl_xor_sync(0xffffffffu, sh, o));
|
||||
sl = __dadd_ru(sl, __shfl_xor_sync(0xffffffffu, sl, o));
|
||||
}
|
||||
if (lane == 0) { s_hn[warp][0] = sh; s_hn[warp][1] = sl; }
|
||||
if (t == 0) { s_n = 0; s_bad = 0; }
|
||||
if (t < 8) s_cnt[t] = 0;
|
||||
__syncthreads();
|
||||
if (t == 0) {
|
||||
double a = 0.0, c = 0.0;
|
||||
for (int hh = 0; hh < IDX_HEADS; ++hh) {
|
||||
double eh = 0.0, el = 0.0;
|
||||
for (int w = 0; w < 4; ++w) { eh = __dadd_ru(eh, s_hn[4 * hh + w][0]); el = __dadd_ru(el, s_hn[4 * hh + w][1]); }
|
||||
a = __dadd_ru(a, __dsqrt_ru(eh));
|
||||
c = __dadd_ru(c, __dsqrt_ru(el));
|
||||
}
|
||||
double P = 0.0, Q = 0.0;
|
||||
tp::coefficients(a, c, P, Q); // RN in double: far inside kSlack
|
||||
s_P = __double2float_ru(P);
|
||||
s_Q = __double2float_ru(Q);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
const float P = s_P, Q = s_Q;
|
||||
const int64_t nb = n_bid + 1;
|
||||
const int w_tail = (int) (n_kv - n_bid * R);
|
||||
const float* sc = scores + qi * max_blocks;
|
||||
bool full = !(P <= tp::kMaxEps && Q <= tp::kMaxEps);
|
||||
float F = -tp_inf();
|
||||
if (!full && nb > tp::kCandCap) { // 2. the sample
|
||||
int64_t stride = tp::kSampleStride;
|
||||
while ((nb + stride - 1) / stride > tp::kCandCap) stride *= 2;
|
||||
const int ns = (int) ((nb + stride - 1) / stride);
|
||||
for (int i = t; i < ns; i += TPB_T) {
|
||||
const int64_t b = (int64_t) i * stride;
|
||||
uint32_t k = 0u; // no bound, or no cells: the lowest key
|
||||
if (b < n_bid) {
|
||||
float lo, hi, e;
|
||||
if (tp_interval(sc[b], kn[b], P, Q, lo, hi, e)) k = order_key(lo);
|
||||
} else if (w_tail > 0) {
|
||||
k = order_key(sc[b]);
|
||||
}
|
||||
c_key[i] = k;
|
||||
}
|
||||
__syncthreads();
|
||||
const int ks = (int) (2 * ((width + R * stride - 1) / (R * stride)) + 16);
|
||||
const uint32_t fk = tp_select(c_key, c_blk, ns, ks < ns ? ks : ns, true, n_bid, w_tail, hist, s_misc);
|
||||
if (fk != 0u) F = tp_key_float(fk);
|
||||
}
|
||||
int wf = 0;
|
||||
float emax = 0.0f;
|
||||
bool bad = false;
|
||||
if (!full) { // 3. one read of every block
|
||||
for (int64_t b = t; b < nb; b += TPB_T) {
|
||||
float lo, hi, e = 0.0f, s;
|
||||
int w;
|
||||
if (b < n_bid) {
|
||||
w = R;
|
||||
s = sc[b];
|
||||
if (!tp_interval(s, kn[b], P, Q, lo, hi, e)) { bad = true; continue; }
|
||||
} else {
|
||||
w = w_tail;
|
||||
if (w == 0) continue;
|
||||
s = lo = hi = sc[b]; // the tail kernel's score: today's
|
||||
}
|
||||
if (lo >= F) wf += w;
|
||||
if (hi >= F) {
|
||||
const int pos = atomicAdd(&s_n, 1);
|
||||
if (pos < tp::kCandCap) { c_key[pos] = order_key(s); c_blk[pos] = (int32_t) b; }
|
||||
emax = fmaxf(emax, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
wf += __shfl_xor_sync(0xffffffffu, wf, o);
|
||||
emax = fmaxf(emax, __shfl_xor_sync(0xffffffffu, emax, o));
|
||||
}
|
||||
bad = __any_sync(0xffffffffu, bad);
|
||||
if (lane == 0) { s_wf[warp] = wf; s_e[warp] = emax; if (bad) s_bad = 1; }
|
||||
__syncthreads();
|
||||
if (t == 0) {
|
||||
int W = 0;
|
||||
float E = 0.0f;
|
||||
for (int w = 0; w < TPB_T / 32; ++w) { W += s_wf[w]; E = fmaxf(E, s_e[w]); }
|
||||
s_misc[0] = W;
|
||||
s_E = E;
|
||||
}
|
||||
__syncthreads();
|
||||
const int n = s_n;
|
||||
full = full || s_bad != 0 || n > tp::kCandCap || s_misc[0] < width;
|
||||
if (!full) {
|
||||
// 4. T1, then 5. the window [lo_b, hi_b] that holds today's threshold
|
||||
const uint32_t t1k = tp_select(c_key, c_blk, n, (int) width, false, n_bid, w_tail, hist, s_misc);
|
||||
const float T1 = tp_key_float(t1k), E = s_E;
|
||||
const float lo_b = fmaxf(__fsub_rd(T1, E), F), hi_b = __fadd_ru(T1, E);
|
||||
for (int i = t; i < n; i += TPB_T) {
|
||||
const int64_t b = c_blk[i];
|
||||
if (b >= n_bid) continue; // the tail block holds today's score already
|
||||
float lo, hi, e;
|
||||
tp_interval(tp_key_float(c_key[i]), kn[b], P, Q, lo, hi, e); // bounded in 3: the same result
|
||||
if (hi >= lo_b && lo <= hi_b) {
|
||||
const int r = (int) (b & 7);
|
||||
const int p = atomicAdd(&s_cnt[r], 1);
|
||||
if (p < tp::kResCap) ids[(qi * 8 + r) * tp::kResCap + p] = (int32_t) b;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (t == 0) {
|
||||
int tot = 0, mx = 0;
|
||||
for (int r = 0; r < 8; ++r) { tot += s_cnt[r]; mx = max(mx, s_cnt[r]); }
|
||||
if (mx > tp::kResCap) full = true;
|
||||
atomicAdd(&g_tp_stats[tp::kStQueries], 1ull);
|
||||
atomicAdd(&g_tp_stats[tp::kStBlocks], (unsigned long long) n_bid);
|
||||
if (full) {
|
||||
mode[qi] = 2;
|
||||
if (atomicExch(&tiles[qi / TC_QT], 1) == 0) atomicAdd(&g_tp_stats[tp::kStFullTiles], 1ull);
|
||||
atomicAdd(&g_tp_stats[tp::kStFull], 1ull);
|
||||
} else {
|
||||
for (int r = 0; r < 8; ++r) cnt[qi * 8 + r] = s_cnt[r];
|
||||
qc[qi] = make_float2(P, Q);
|
||||
mode[qi] = 1;
|
||||
atomicAdd(&g_tp_stats[tp::kStBand], (unsigned long long) tot);
|
||||
atomicMax(&g_tp_stats[tp::kStBandMax], (unsigned long long) tot);
|
||||
atomicAdd(&g_tp_stats[tp::kStCand], (unsigned long long) n);
|
||||
atomicMax(&g_tp_stats[tp::kStCandMax], (unsigned long long) n);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The listed blocks scored again exactly as block_scores_tc_kernel scores them: the query in its row of the m16
|
||||
// tile (qi % 16), each block in its column of the n8 tile (b % 8), the same TF32 splits, the same three MMAs per step
|
||||
// in the same order into the same accumulator, the same relu sum. The other rows and columns hold zeros where
|
||||
// today's kernel holds other queries and blocks; an MMA output element depends only on its own row of A, column of B
|
||||
// and accumulator, so the element is today's bit for bit (qsa_two_pass_parity and STRATA_SELECT_TWO_PASS_CHECK
|
||||
// compare every one with today's kernel). Warp w takes tiles w, w + 4, ...: tile j holds the j-th listed block of
|
||||
// every residue.
|
||||
__global__ void __launch_bounds__(128) tp_rescore_kernel(const float* __restrict__ pooled,
|
||||
const float* __restrict__ q_idx, int64_t max_blocks,
|
||||
const int* __restrict__ mode, const int* __restrict__ cnt,
|
||||
const int* __restrict__ ids, const int* __restrict__ tiles,
|
||||
float* __restrict__ out) {
|
||||
__shared__ __align__(16) float sq[IDX_HEADS * IDX_DIM];
|
||||
__shared__ int s_cnt[8];
|
||||
const int64_t qi = blockIdx.x;
|
||||
if (mode[qi] != 1 || tiles[qi / TC_QT] != 0) return; // nothing listed, or the full kernel takes the tile
|
||||
const int t = threadIdx.x, lane = t & 31, warp = t >> 5, gid = lane >> 2, tig = lane & 3;
|
||||
for (int i = t; i < IDX_HEADS * IDX_DIM; i += 128) sq[i] = q_idx[qi * IDX_HEADS * IDX_DIM + i];
|
||||
if (t < 8) s_cnt[t] = cnt[qi * 8 + t];
|
||||
__syncthreads();
|
||||
int ntile = 0;
|
||||
for (int r = 0; r < 8; ++r) ntile = max(ntile, s_cnt[r]);
|
||||
const int row = (int) (qi % TC_QT);
|
||||
const bool mine0 = gid == (row & 7) && row < 8, mine1 = gid == (row & 7) && row >= 8;
|
||||
const int* lst = ids + qi * 8 * tp::kResCap;
|
||||
for (int j = warp; j < ntile; j += 4) {
|
||||
const int bcol = j < s_cnt[gid] ? lst[gid * tp::kResCap + j] : -1; // this lane's column: gid
|
||||
const float* kr = pooled + (int64_t) (bcol >= 0 ? bcol : 0) * IDX_DIM;
|
||||
float acc[IDX_HEADS][4];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < IDX_HEADS; ++h) acc[h][0] = acc[h][1] = acc[h][2] = acc[h][3] = 0.f;
|
||||
#pragma unroll 4
|
||||
for (int k0 = 0; k0 < IDX_DIM; k0 += 8) {
|
||||
const float kx0 = bcol >= 0 ? kr[k0 + tig] : 0.f, kx1 = bcol >= 0 ? kr[k0 + tig + 4] : 0.f;
|
||||
uint32_t bh[2], bl[2];
|
||||
bh[0] = tf32_hi(kx0);
|
||||
bh[1] = tf32_hi(kx1);
|
||||
bl[0] = tf32_hi(kx0 - __uint_as_float(bh[0]));
|
||||
bl[1] = tf32_hi(kx1 - __uint_as_float(bh[1]));
|
||||
#pragma unroll
|
||||
for (int h = 0; h < IDX_HEADS; ++h) {
|
||||
const float* qa = sq + h * IDX_DIM + k0 + tig;
|
||||
const float x0 = mine0 ? qa[0] : 0.f, x1 = mine1 ? qa[0] : 0.f;
|
||||
const float x2 = mine0 ? qa[4] : 0.f, x3 = mine1 ? qa[4] : 0.f;
|
||||
uint32_t ah[4], al[4];
|
||||
ah[0] = tf32_hi(x0); ah[1] = tf32_hi(x1); ah[2] = tf32_hi(x2); ah[3] = tf32_hi(x3);
|
||||
al[0] = tf32_hi(x0 - __uint_as_float(ah[0]));
|
||||
al[1] = tf32_hi(x1 - __uint_as_float(ah[1]));
|
||||
al[2] = tf32_hi(x2 - __uint_as_float(ah[2]));
|
||||
al[3] = tf32_hi(x3 - __uint_as_float(ah[3]));
|
||||
mma_tf32(acc[h], al, bh);
|
||||
mma_tf32(acc[h], ah, bl);
|
||||
mma_tf32(acc[h], ah, bh);
|
||||
}
|
||||
}
|
||||
if (gid == (row & 7)) { // row `row`, columns 2 tig and 2 tig + 1
|
||||
#pragma unroll
|
||||
for (int jj = 0; jj < 2; ++jj) {
|
||||
const int c = 2 * tig + jj;
|
||||
if (j >= s_cnt[c]) continue;
|
||||
const int64_t b = lst[c * tp::kResCap + j];
|
||||
float score = 0.0f;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < IDX_HEADS; ++h) {
|
||||
const float d = row >= 8 ? (jj ? acc[h][3] : acc[h][2]) : (jj ? acc[h][1] : acc[h][0]);
|
||||
score += d > 0.0f ? d : 0.0f;
|
||||
}
|
||||
out[qi * max_blocks + b] = score;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A flagged 16-query tile scored again in full: block_scores_tc_kernel's body as it is above, the CTA taking the tile
|
||||
// (blockIdx.y) only when the band kernel flagged it and striding over its 128-block ranges. Each 32-block tile is
|
||||
// computed exactly as today's kernel computes it (the same shared-memory tiles, warps, MMAs and sums), so every
|
||||
// score of the tile is today's.
|
||||
__global__ void __launch_bounds__(128) tp_full_kernel(const float* __restrict__ pooled,
|
||||
const float* __restrict__ q_idx,
|
||||
const int32_t* __restrict__ steps, int64_t nq,
|
||||
int64_t max_blocks, int64_t reach,
|
||||
const int* __restrict__ tiles, float* __restrict__ out) {
|
||||
if (tiles[blockIdx.y] == 0) return;
|
||||
extern __shared__ __align__(16) float sm[];
|
||||
float* sQ = sm; // [TC_QT][TC_QS]
|
||||
float* sK = sm + TC_QT * TC_QS; // [TC_NB][TC_KS]
|
||||
__shared__ int s_nbid[TC_QT];
|
||||
const int t = threadIdx.x, lane = t & 31, warp = t >> 5, gid = lane >> 2, tig = lane & 3;
|
||||
const int64_t q0 = (int64_t) blockIdx.y * TC_QT;
|
||||
for (int i = t; i < TC_QT * IDX_HEADS * IDX_DIM / 4; i += 128) {
|
||||
const int r = i / (IDX_HEADS * IDX_DIM / 4), c = i % (IDX_HEADS * IDX_DIM / 4);
|
||||
float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
if (q0 + r < nq) v = reinterpret_cast<const float4*>(q_idx + (q0 + r) * IDX_HEADS * IDX_DIM)[c];
|
||||
*reinterpret_cast<float4*>(sQ + r * TC_QS + c * 4) = v;
|
||||
}
|
||||
if (t < TC_QT) s_nbid[t] = q0 + t < nq ? steps[(q0 + t) * kStepCount + kStepNBid] : 0;
|
||||
int lo_nbid = 0x7fffffff, hi_nbid = 0;
|
||||
__syncthreads();
|
||||
for (int i = 0; i < TC_QT; ++i) {
|
||||
if (q0 + i >= nq) break;
|
||||
lo_nbid = min(lo_nbid, s_nbid[i]);
|
||||
hi_nbid = max(hi_nbid, s_nbid[i]);
|
||||
}
|
||||
(void) lo_nbid;
|
||||
for (int64_t x = blockIdx.x; x * TC_NB * TC_ITER < reach; x += gridDim.x) {
|
||||
for (int it = 0; it < TC_ITER; ++it) {
|
||||
const int64_t b0 = (x * TC_ITER + it) * TC_NB;
|
||||
if (b0 >= reach || b0 >= hi_nbid) break;
|
||||
__syncthreads();
|
||||
for (int i = t; i < TC_NB * IDX_DIM / 4; i += 128) {
|
||||
const int r = i / (IDX_DIM / 4), c = i % (IDX_DIM / 4);
|
||||
float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
if (b0 + r < hi_nbid) v = reinterpret_cast<const float4*>(pooled + (b0 + r) * IDX_DIM)[c];
|
||||
*reinterpret_cast<float4*>(sK + r * TC_KS + c * 4) = v;
|
||||
}
|
||||
__syncthreads();
|
||||
float acc[IDX_HEADS][4];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < IDX_HEADS; ++h) acc[h][0] = acc[h][1] = acc[h][2] = acc[h][3] = 0.f;
|
||||
const float* kr = sK + (warp * 8 + gid) * TC_KS;
|
||||
#pragma unroll 4
|
||||
for (int k0 = 0; k0 < IDX_DIM; k0 += 8) {
|
||||
const float kx0 = kr[k0 + tig], kx1 = kr[k0 + tig + 4];
|
||||
uint32_t bh[2], bl[2];
|
||||
bh[0] = tf32_hi(kx0);
|
||||
bh[1] = tf32_hi(kx1);
|
||||
bl[0] = tf32_hi(kx0 - __uint_as_float(bh[0]));
|
||||
bl[1] = tf32_hi(kx1 - __uint_as_float(bh[1]));
|
||||
#pragma unroll
|
||||
for (int h = 0; h < IDX_HEADS; ++h) {
|
||||
const float* qa = sQ + h * IDX_DIM + k0 + tig;
|
||||
const float x0 = qa[gid * TC_QS], x1 = qa[(gid + 8) * TC_QS];
|
||||
const float x2 = qa[gid * TC_QS + 4], x3 = qa[(gid + 8) * TC_QS + 4];
|
||||
uint32_t ah[4], al[4];
|
||||
ah[0] = tf32_hi(x0); ah[1] = tf32_hi(x1); ah[2] = tf32_hi(x2); ah[3] = tf32_hi(x3);
|
||||
al[0] = tf32_hi(x0 - __uint_as_float(ah[0]));
|
||||
al[1] = tf32_hi(x1 - __uint_as_float(ah[1]));
|
||||
al[2] = tf32_hi(x2 - __uint_as_float(ah[2]));
|
||||
al[3] = tf32_hi(x3 - __uint_as_float(ah[3]));
|
||||
mma_tf32(acc[h], al, bh);
|
||||
mma_tf32(acc[h], ah, bl);
|
||||
mma_tf32(acc[h], ah, bh);
|
||||
}
|
||||
}
|
||||
const int64_t bc = b0 + warp * 8 + 2 * tig;
|
||||
#pragma unroll
|
||||
for (int half = 0; half < 2; ++half) {
|
||||
const int qr = gid + half * 8;
|
||||
const int64_t qi = q0 + qr;
|
||||
if (qi >= nq) continue;
|
||||
const int nbq = s_nbid[qr];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 2; ++j) {
|
||||
const int64_t b = bc + j;
|
||||
if (b >= nbq || b >= max_blocks) continue;
|
||||
float score = 0.0f;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < IDX_HEADS; ++h) {
|
||||
const float d = acc[h][half * 2 + j];
|
||||
score += d > 0.0f ? d : 0.0f;
|
||||
}
|
||||
out[qi * max_blocks + b] = score;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// STRATA_SELECT_TWO_PASS_CHECK: one query against today's row and ids. Ids: every entry qsa_block_topk wrote. Scored-
|
||||
// again blocks (listed, the tail, or every block of a full query): bitwise today's score. Kept blocks: |S1 - S3| <= eps.
|
||||
__global__ void __launch_bounds__(256) tp_check_kernel(const float* __restrict__ scores, const int32_t* __restrict__ ids,
|
||||
const float* __restrict__ ref, const int32_t* __restrict__ ref_ids,
|
||||
const int32_t* __restrict__ steps, int64_t max_blocks,
|
||||
int64_t cap, const float2* __restrict__ kn,
|
||||
const float2* __restrict__ qc, const int* __restrict__ mode,
|
||||
const int* __restrict__ cnt, const int* __restrict__ lists,
|
||||
const int* __restrict__ tiles) {
|
||||
extern __shared__ uint32_t tp_mark[];
|
||||
__shared__ unsigned long long s_acc[6];
|
||||
__shared__ unsigned int s_ratio;
|
||||
const int t = threadIdx.x;
|
||||
const int64_t qi = blockIdx.x;
|
||||
const int32_t* st = steps + qi * kStepCount;
|
||||
const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth];
|
||||
if (t < 6) s_acc[t] = 0ull;
|
||||
if (t == 0) s_ratio = 0u;
|
||||
const int64_t m = n_kv < width ? n_kv : width;
|
||||
unsigned long long dif = 0;
|
||||
for (int64_t j = t; j < m; j += 256) dif += ids[qi * cap + j] != ref_ids[qi * cap + j];
|
||||
const bool sparse = n_kv > width;
|
||||
const int md = !sparse ? 0 : tiles[qi / TC_QT] != 0 ? 2 : mode[qi]; // a flagged tile: every block scored again
|
||||
const int64_t nb = n_bid + 1, words = (nb + 31) / 32;
|
||||
if (sparse) for (int64_t i = t; i < words; i += 256) tp_mark[i] = 0u;
|
||||
__syncthreads();
|
||||
atomicAdd(&s_acc[0], dif);
|
||||
if (sparse && md == 1) {
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
const int c = cnt[qi * 8 + r];
|
||||
for (int j = t; j < c; j += 256) {
|
||||
const int b = lists[(qi * 8 + r) * tp::kResCap + j];
|
||||
atomicOr(&tp_mark[b >> 5], 1u << (b & 31));
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (sparse) {
|
||||
const float* s1 = scores + qi * max_blocks;
|
||||
const float* s3 = ref + qi * max_blocks;
|
||||
const float2 pq = md == 1 ? qc[qi] : make_float2(0.f, 0.f);
|
||||
unsigned long long again = 0, again_dif = 0, kept = 0, kept_out = 0;
|
||||
float ratio = 0.0f;
|
||||
for (int64_t b = t; b < nb; b += 256) {
|
||||
const bool exact = md == 2 || b == n_bid || ((tp_mark[b >> 5] >> (b & 31)) & 1u);
|
||||
if (exact) {
|
||||
++again;
|
||||
again_dif += __float_as_uint(s1[b]) != __float_as_uint(s3[b]);
|
||||
} else {
|
||||
++kept;
|
||||
const float e = __fmaf_ru(pq.x, kn[b].y, __fmaf_ru(pq.y, kn[b].x, tp::kAbs));
|
||||
const double d = fabs((double) s1[b] - (double) s3[b]);
|
||||
kept_out += !(d <= (double) e);
|
||||
const float r = (float) (d / (double) e);
|
||||
ratio = fmaxf(ratio, r == r ? r : tp_inf());
|
||||
}
|
||||
}
|
||||
atomicAdd(&s_acc[1], again);
|
||||
atomicAdd(&s_acc[2], again_dif);
|
||||
atomicAdd(&s_acc[3], kept);
|
||||
atomicAdd(&s_acc[4], kept_out);
|
||||
atomicMax(&s_ratio, __float_as_uint(ratio));
|
||||
}
|
||||
__syncthreads();
|
||||
if (t == 0) {
|
||||
atomicAdd(&g_tp_stats[tp::kChkQueries], 1ull);
|
||||
if (s_acc[0]) atomicAdd(&g_tp_stats[tp::kChkIdQueries], 1ull);
|
||||
atomicAdd(&g_tp_stats[tp::kChkIdEntries], s_acc[0]);
|
||||
atomicAdd(&g_tp_stats[tp::kChkRescored], s_acc[1]);
|
||||
atomicAdd(&g_tp_stats[tp::kChkRescoredDiff], s_acc[2]);
|
||||
atomicAdd(&g_tp_stats[tp::kChkKept], s_acc[3]);
|
||||
atomicAdd(&g_tp_stats[tp::kChkKeptOut], s_acc[4]);
|
||||
atomicMax(&g_tp_stats[tp::kChkRatioMax], (unsigned long long) s_ratio);
|
||||
}
|
||||
}
|
||||
|
||||
// The card runs qsa_block_scores_tc's tensor-core kernel (the reference of the two-pass scores), and the full kernel
|
||||
// takes its shared memory. The same dispatch rule as qsa_block_scores_tc: sm_80+, STRATA_QSA_WARP not set.
|
||||
bool tp_device_ok() {
|
||||
static int ok[64] = {}; // 0 unknown, 1 yes, 2 no
|
||||
int dev = 0;
|
||||
if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return false; }
|
||||
if (ok[dev] == 0) {
|
||||
int major = 0;
|
||||
bool good = cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) == cudaSuccess;
|
||||
const char* w = std::getenv("STRATA_QSA_WARP");
|
||||
const int cc = w && (!std::strcmp(w, "1") || !std::strcmp(w, "select")) ? 7 : strata::cc_major_of(major);
|
||||
const int bytes = (TC_QT * TC_QS + TC_NB * TC_KS) * (int) sizeof(float);
|
||||
good = good && cc >= 8 &&
|
||||
cudaFuncSetAttribute(block_scores_tc_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) ==
|
||||
cudaSuccess &&
|
||||
cudaFuncSetAttribute(tp_full_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) == cudaSuccess;
|
||||
cudaGetLastError();
|
||||
ok[dev] = good ? 1 : 2;
|
||||
}
|
||||
return ok[dev] == 1;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
bool qsa_two_pass_on() {
|
||||
static two_pass::RequestStatus status(std::getenv("STRATA_SELECT_TWO_PASS"));
|
||||
if (!status.requested()) return false;
|
||||
// the card current at this call (a layer split can mix cards): sm_75 or STRATA_QSA_WARP=1|select keep the other
|
||||
// scorer there, so the switch says once that it is not used; checked here too, so a short prompt says it as well
|
||||
if (!tp_device_ok()) (void) status.on("this CUDA card or with STRATA_QSA_WARP");
|
||||
return true;
|
||||
}
|
||||
bool qsa_two_pass_check_on() {
|
||||
static const bool on = [] {
|
||||
const char* v = std::getenv("STRATA_SELECT_TWO_PASS_CHECK");
|
||||
return v && v[0] == '1' && v[1] == '\0';
|
||||
}();
|
||||
return on;
|
||||
}
|
||||
int64_t qsa_two_pass_min_blocks() {
|
||||
static const int64_t v = [] {
|
||||
const char* e = std::getenv("STRATA_SELECT_TWO_PASS_MIN");
|
||||
const long long x = e ? std::atoll(e) : (long long) two_pass::kMinBlocksDefault;
|
||||
return (int64_t) (x < 0 ? 0 : x);
|
||||
}();
|
||||
return v;
|
||||
}
|
||||
uint64_t qsa_two_pass_scratch_bytes(int64_t nq_cap, int64_t max_blocks) { return two_pass::layout(nq_cap, max_blocks).total; }
|
||||
|
||||
bool qsa_two_pass_keys(const float* pooled, int64_t n_blocks, int64_t nq_cap, int64_t max_blocks, void* scratch,
|
||||
uint64_t scratch_bytes, void* stream) {
|
||||
if (!qsa_two_pass_on() || !tp_device_ok() || n_blocks > max_blocks || nq_cap <= 0 || nq_cap > 4096) return false;
|
||||
const two_pass::Layout L = two_pass::layout(nq_cap, max_blocks);
|
||||
if (scratch == nullptr || scratch_bytes < L.total || (reinterpret_cast<uintptr_t>(scratch) & 15u) != 0) return false;
|
||||
if (n_blocks <= 0) return true;
|
||||
float2* kn = reinterpret_cast<float2*>(static_cast<uint8_t*>(scratch) + L.kn);
|
||||
tp_keys_kernel<<<(unsigned) ((n_blocks + 7) / 8), 256, 0, (cudaStream_t) stream>>>(pooled, n_blocks, kn);
|
||||
const cudaError_t e = cudaGetLastError();
|
||||
if (e != cudaSuccess) { std::fprintf(stderr, "qsa two-pass keys: %s\n", cudaGetErrorString(e)); std::exit(1); }
|
||||
return true;
|
||||
}
|
||||
|
||||
bool qsa_block_scores_two_pass(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps,
|
||||
int64_t nq, int64_t nq_cap, int64_t max_blocks, const QsaShapes& s, float* scores,
|
||||
void* scratch, uint64_t scratch_bytes, void* stream, int64_t active_blocks) {
|
||||
if (nq <= 0) return true;
|
||||
if (!qsa_two_pass_on() || s.idx_dim != IDX_DIM || s.idx_n_head != IDX_HEADS || s.idx_block != R ||
|
||||
nq > nq_cap || nq_cap > 4096 || !tp_device_ok())
|
||||
return false;
|
||||
const two_pass::Layout L = two_pass::layout(nq_cap, max_blocks);
|
||||
if (scratch == nullptr || scratch_bytes < L.total || (reinterpret_cast<uintptr_t>(scratch) & 15u) != 0) return false;
|
||||
uint8_t* base = static_cast<uint8_t*>(scratch);
|
||||
const float2* kn = reinterpret_cast<const float2*>(base + L.kn);
|
||||
float2* qc = reinterpret_cast<float2*>(base + L.qc);
|
||||
int* mode = reinterpret_cast<int*>(base + L.mode);
|
||||
int* cnt = reinterpret_cast<int*>(base + L.cnt);
|
||||
int* ids = reinterpret_cast<int*>(base + L.ids);
|
||||
int* tiles = reinterpret_cast<int*>(base + L.tiles);
|
||||
const cudaStream_t st = (cudaStream_t) stream;
|
||||
const int64_t reach = active_blocks > 0 && active_blocks < max_blocks ? active_blocks : max_blocks;
|
||||
static std::atomic<bool> said{false};
|
||||
if (!said.exchange(true))
|
||||
std::fprintf(stderr, "qsa select: two-pass scores on (STRATA_SELECT_TWO_PASS=1): calls whose last query has %lld+ "
|
||||
"blocks score every block in single-pass TF32 and again in 3xTF32 where the bound reaches the "
|
||||
"cut-off (the same cells as the 3xTF32 scorer)\n", (long long) qsa_two_pass_min_blocks());
|
||||
const unsigned qtiles = (unsigned) ((nq + TC_QT - 1) / TC_QT);
|
||||
const int64_t per = (int64_t) TC_NB * TP1_ITER;
|
||||
tp_score1_kernel<<<dim3(qtiles, (unsigned) ((reach + per - 1) / per)), 128, 0, st>>>(pooled, q_idx, steps, nq,
|
||||
max_blocks, reach, scores, tiles);
|
||||
block_scores_tail_kernel<<<(unsigned) nq, 32, 0, st>>>(dead, q_idx, steps, max_blocks, scores);
|
||||
tp_band_kernel<<<(unsigned) nq, TPB_T, (size_t) two_pass::kCandCap * 8, st>>>(scores, q_idx, steps, max_blocks, kn, qc,
|
||||
mode, cnt, ids, tiles);
|
||||
tp_rescore_kernel<<<(unsigned) nq, 128, 0, st>>>(pooled, q_idx, max_blocks, mode, cnt, ids, tiles, scores);
|
||||
const int64_t ranges = (reach + (int64_t) TC_NB * TC_ITER - 1) / ((int64_t) TC_NB * TC_ITER);
|
||||
const int bytes = (TC_QT * TC_QS + TC_NB * TC_KS) * (int) sizeof(float);
|
||||
tp_full_kernel<<<dim3((unsigned) (ranges < 32 ? ranges : 32), qtiles), 128, bytes, st>>>(pooled, q_idx, steps, nq,
|
||||
max_blocks, reach, tiles,
|
||||
scores);
|
||||
const cudaError_t e = cudaGetLastError();
|
||||
if (e != cudaSuccess) { std::fprintf(stderr, "qsa two-pass scores: %s\n", cudaGetErrorString(e)); std::exit(1); }
|
||||
return true;
|
||||
}
|
||||
|
||||
void qsa_two_pass_check(const float* scores, const int32_t* ids, const float* ref_scores, const int32_t* ref_ids,
|
||||
const int32_t* steps, int64_t nq, int64_t nq_cap, int64_t max_blocks, int64_t cap,
|
||||
void* scratch, void* stream) {
|
||||
if (nq <= 0 || scratch == nullptr) return;
|
||||
const size_t mark = (size_t) ((max_blocks + 31) / 32) * 4;
|
||||
if (mark > 48 * 1024) {
|
||||
static std::atomic<bool> said{false};
|
||||
if (!said.exchange(true)) std::fprintf(stderr, "qsa two-pass check: %lld blocks are too many to check\n", (long long) max_blocks);
|
||||
return;
|
||||
}
|
||||
const two_pass::Layout L = two_pass::layout(nq_cap, max_blocks);
|
||||
const uint8_t* base = static_cast<const uint8_t*>(scratch);
|
||||
tp_check_kernel<<<(unsigned) nq, 256, mark, (cudaStream_t) stream>>>(
|
||||
scores, ids, ref_scores, ref_ids, steps, max_blocks, cap, reinterpret_cast<const float2*>(base + L.kn),
|
||||
reinterpret_cast<const float2*>(base + L.qc), reinterpret_cast<const int*>(base + L.mode),
|
||||
reinterpret_cast<const int*>(base + L.cnt), reinterpret_cast<const int*>(base + L.ids),
|
||||
reinterpret_cast<const int*>(base + L.tiles));
|
||||
const cudaError_t e = cudaGetLastError();
|
||||
if (e != cudaSuccess) { std::fprintf(stderr, "qsa two-pass check: %s\n", cudaGetErrorString(e)); std::exit(1); }
|
||||
}
|
||||
|
||||
bool qsa_two_pass_stats(unsigned long long* out, bool reset, void* stream) {
|
||||
const cudaStream_t st = (cudaStream_t) stream;
|
||||
void* p = nullptr;
|
||||
if (cudaGetSymbolAddress(&p, g_tp_stats) != cudaSuccess) { cudaGetLastError(); return false; }
|
||||
const size_t bytes = sizeof(unsigned long long) * two_pass::kStatCount;
|
||||
if (cudaMemcpyAsync(out, p, bytes, cudaMemcpyDeviceToHost, st) != cudaSuccess ||
|
||||
(reset && cudaMemsetAsync(p, 0, bytes, st) != cudaSuccess) || cudaStreamSynchronize(st) != cudaSuccess) {
|
||||
cudaGetLastError();
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
#else // HIP: the two-pass scorer is CUDA-only (the gfx12 WMMA scorer has another arithmetic); nothing runs
|
||||
bool qsa_two_pass_on() { // false; a set switch says once that it is not used here
|
||||
static two_pass::RequestStatus status(std::getenv("STRATA_SELECT_TWO_PASS"));
|
||||
return status.on("HIP");
|
||||
}
|
||||
bool qsa_two_pass_check_on() { return false; }
|
||||
int64_t qsa_two_pass_min_blocks() { return two_pass::kMinBlocksDefault; }
|
||||
uint64_t qsa_two_pass_scratch_bytes(int64_t nq_cap, int64_t max_blocks) { return two_pass::layout(nq_cap, max_blocks).total; }
|
||||
bool qsa_two_pass_keys(const float*, int64_t, int64_t, int64_t, void*, uint64_t, void*) { return false; }
|
||||
bool qsa_block_scores_two_pass(const float*, const float*, const float*, const int32_t*, int64_t, int64_t, int64_t,
|
||||
const QsaShapes&, float*, void*, uint64_t, void*, int64_t) { return false; }
|
||||
void qsa_two_pass_check(const float*, const int32_t*, const float*, const int32_t*, const int32_t*, int64_t, int64_t,
|
||||
int64_t, int64_t, void*, void*) {}
|
||||
bool qsa_two_pass_stats(unsigned long long*, bool, void*) { return false; }
|
||||
#endif
|
||||
|
||||
} // namespace strata::kernels
|
||||
@@ -0,0 +1,224 @@
|
||||
// STRATA_SELECT_TWO_PASS (qsa_two_pass.hpp) against today's prompt scorer, GPU, synthetic, no model: for prompt
|
||||
// chunks of 256 queries at depths from 2K to 256K and several kinds of keys (normal, indexer-like, wide magnitudes,
|
||||
// small integers with exact ties, near-duplicate and duplicate keys), qsa_block_topk over the two-pass row must return
|
||||
// exactly the ids it returns over qsa_block_scores_tc's row, every scored-again block must hold today's score bit for
|
||||
// bit, and every kept single-pass score must lie inside its bound (qsa_two_pass_check's counters). Then, for the
|
||||
// larger cases, the time of each scorer (one synthetic run, not a benchmark). Needs CUDA sm_80+ (77 = skipped
|
||||
// otherwise, and on HIP).
|
||||
#include "strata/kernels/qsa.hpp"
|
||||
#include "strata/kernels/qsa_select.hpp"
|
||||
#include "strata/kernels/qsa_two_pass.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
namespace k = strata::kernels;
|
||||
namespace tp = strata::kernels::two_pass;
|
||||
|
||||
namespace {
|
||||
void ck(cudaError_t e) {
|
||||
if (e != cudaSuccess) {
|
||||
std::fprintf(stderr, "CUDA: %s\n", cudaGetErrorString(e));
|
||||
std::exit(2);
|
||||
}
|
||||
}
|
||||
|
||||
enum Kind { kNormal, kIndexer, kWide, kSmall, kNear, kDup, kKinds };
|
||||
const char* kKindName[kKinds] = {"normal", "indexer-like", "wide magnitudes", "small integers", "near-duplicate keys",
|
||||
"duplicate keys"};
|
||||
|
||||
struct Case { int64_t ctx, queries; Kind kind; bool time; };
|
||||
|
||||
// stride: the score row's stride in blocks; 0 = the case's own (ctx / 4 + 2), else the engine's (--max-context 262144:
|
||||
// 65,538, so qsa_block_topk takes the engine's dispatch on this card)
|
||||
bool run(const Case& c, int64_t stride) {
|
||||
const auto s = k::qsa_real_shapes();
|
||||
const int64_t D = s.idx_dim, H = s.idx_n_head, own = c.ctx / s.idx_block + 2;
|
||||
const int64_t max_blocks = stride > own ? stride : own;
|
||||
const int64_t nq = c.queries, nq_cap = 256, cap = k::qsa_selection_width(k::kTopkMaxCells, s);
|
||||
std::mt19937 rng(77001 + (unsigned) c.ctx + 13 * (unsigned) c.queries + 1009 * (unsigned) c.kind);
|
||||
std::normal_distribution<float> nd(0.f, 1.f);
|
||||
std::uniform_int_distribution<int> small(-3, 3), expo(-16, 16), coin(0, 9);
|
||||
std::vector<float> pooled((size_t) (max_blocks * D)), dead((size_t) D), q((size_t) (nq * H * D));
|
||||
auto unit = [&](float* v, int n, float norm) {
|
||||
double ss = 0;
|
||||
for (int i = 0; i < n; ++i) ss += (double) v[i] * v[i];
|
||||
const float f = ss > 0 ? (float) (norm / std::sqrt(ss)) : 0.f;
|
||||
for (int i = 0; i < n; ++i) v[i] *= f;
|
||||
};
|
||||
for (int64_t b = 0; b < max_blocks; ++b) {
|
||||
float* kb = pooled.data() + b * D;
|
||||
if (c.kind == kIndexer) { // the mean of 4 RMS-normed cell keys, as the indexer pools them
|
||||
for (int i = 0; i < D; ++i) kb[i] = 0.f;
|
||||
for (int cell = 0; cell < 4; ++cell) {
|
||||
float ck_[128];
|
||||
for (int i = 0; i < D; ++i) ck_[i] = nd(rng);
|
||||
unit(ck_, (int) D, std::sqrt((float) D));
|
||||
for (int i = 0; i < D; ++i) kb[i] += 0.25f * ck_[i];
|
||||
}
|
||||
} else if (c.kind == kNear && b > 0 && coin(rng) < 5) { // half the keys a hair from the one before
|
||||
for (int i = 0; i < D; ++i) kb[i] = kb[i - D] * (1.0f + 1e-6f * nd(rng));
|
||||
} else if (c.kind == kDup && b >= 8 && coin(rng) < 5) { // half the keys copies of an earlier one: exact ties
|
||||
const int64_t src = (int64_t) (rng() % (uint32_t) b);
|
||||
std::memcpy(kb, pooled.data() + src * D, (size_t) D * 4);
|
||||
} else {
|
||||
for (int i = 0; i < D; ++i) {
|
||||
if (c.kind == kWide) kb[i] = coin(rng) == 0 ? 0.f : std::ldexp(nd(rng), expo(rng));
|
||||
else if (c.kind == kSmall) kb[i] = (float) small(rng);
|
||||
else kb[i] = nd(rng);
|
||||
}
|
||||
}
|
||||
}
|
||||
for (auto& v : dead) v = nd(rng);
|
||||
for (int64_t i = 0; i < nq; ++i)
|
||||
for (int64_t h = 0; h < H; ++h) {
|
||||
float* qh = q.data() + (i * H + h) * D;
|
||||
if (c.kind == kIndexer) { // RMS-normed, leaning towards a few keys (a peaked score row)
|
||||
for (int j = 0; j < D; ++j) qh[j] = nd(rng);
|
||||
for (int r = 0; r < 3; ++r) {
|
||||
const float* kb = pooled.data() + (int64_t) (rng() % (uint32_t) (max_blocks - 2)) * D;
|
||||
for (int j = 0; j < D; ++j) qh[j] += 0.5f * kb[j];
|
||||
}
|
||||
unit(qh, (int) D, std::sqrt((float) D));
|
||||
} else {
|
||||
for (int j = 0; j < D; ++j) {
|
||||
if (c.kind == kWide) qh[j] = coin(rng) == 0 ? 0.f : std::ldexp(nd(rng), expo(rng));
|
||||
else if (c.kind == kSmall) qh[j] = (float) small(rng);
|
||||
else qh[j] = nd(rng);
|
||||
}
|
||||
}
|
||||
}
|
||||
std::vector<int32_t> steps((size_t) (nq * k::kStepCount));
|
||||
for (int64_t i = 0; i < nq; ++i) {
|
||||
int32_t* st = steps.data() + i * k::kStepCount;
|
||||
const int64_t pos = c.ctx - nq + i;
|
||||
st[k::kStepPos] = (int32_t) pos;
|
||||
st[k::kStepNKv] = (int32_t) (pos + 1);
|
||||
st[k::kStepNBid] = (int32_t) ((pos + 1) / s.idx_block);
|
||||
st[k::kStepWidth] = (int32_t) k::qsa_selection_width(pos + 1, s);
|
||||
}
|
||||
const int64_t last_nbid = steps[(size_t) ((nq - 1) * k::kStepCount + k::kStepNBid)], active = last_nbid + 1;
|
||||
float *dp, *dd, *dq, *ref, *two;
|
||||
int32_t *dt, *ref_ids, *two_ids;
|
||||
void* scratch;
|
||||
const uint64_t sbytes = k::qsa_two_pass_scratch_bytes(nq_cap, max_blocks);
|
||||
const size_t n_sc = (size_t) (nq * max_blocks), n_id = (size_t) (nq * cap);
|
||||
ck(cudaMalloc(&dp, pooled.size() * 4)); ck(cudaMalloc(&dd, dead.size() * 4)); ck(cudaMalloc(&dq, q.size() * 4));
|
||||
ck(cudaMalloc(&dt, steps.size() * 4)); ck(cudaMalloc(&ref, n_sc * 4)); ck(cudaMalloc(&two, n_sc * 4));
|
||||
ck(cudaMalloc(&ref_ids, n_id * 4)); ck(cudaMalloc(&two_ids, n_id * 4)); ck(cudaMalloc(&scratch, sbytes));
|
||||
ck(cudaMemcpy(dp, pooled.data(), pooled.size() * 4, cudaMemcpyHostToDevice));
|
||||
ck(cudaMemcpy(dd, dead.data(), dead.size() * 4, cudaMemcpyHostToDevice));
|
||||
ck(cudaMemcpy(dq, q.data(), q.size() * 4, cudaMemcpyHostToDevice));
|
||||
ck(cudaMemcpy(dt, steps.data(), steps.size() * 4, cudaMemcpyHostToDevice));
|
||||
ck(cudaMemset(ref, 0x7f, n_sc * 4)); ck(cudaMemset(two, 0x3e, n_sc * 4));
|
||||
ck(cudaMemset(ref_ids, 0xff, n_id * 4)); ck(cudaMemset(two_ids, 0xee, n_id * 4));
|
||||
if (!k::qsa_block_scores_tc(dp, dd, dq, dt, nq, max_blocks, s, ref, nullptr, active)) {
|
||||
std::puts("SKIP: no tensor-core scorer on this card");
|
||||
std::exit(77);
|
||||
}
|
||||
k::qsa_block_topk(ref, dt, nq, max_blocks, cap, s, ref_ids, nullptr, active);
|
||||
unsigned long long st0[tp::kStatCount];
|
||||
k::qsa_two_pass_stats(st0, true, nullptr); // zero the counters
|
||||
if (!k::qsa_two_pass_keys(dp, last_nbid, nq_cap, max_blocks, scratch, sbytes, nullptr) ||
|
||||
!k::qsa_block_scores_two_pass(dp, dd, dq, dt, nq, nq_cap, max_blocks, s, two, scratch, sbytes, nullptr, active)) {
|
||||
std::fprintf(stderr, "FAIL: the two-pass scorer did not run (ctx %lld)\n", (long long) c.ctx);
|
||||
return false;
|
||||
}
|
||||
k::qsa_block_topk(two, dt, nq, max_blocks, cap, s, two_ids, nullptr, active);
|
||||
k::qsa_two_pass_check(two, two_ids, ref, ref_ids, dt, nq, nq_cap, max_blocks, cap, scratch, nullptr);
|
||||
ck(cudaDeviceSynchronize());
|
||||
unsigned long long st[tp::kStatCount];
|
||||
if (!k::qsa_two_pass_stats(st, true, nullptr)) { std::fprintf(stderr, "FAIL: counters\n"); return false; }
|
||||
// the ids on the host too, every entry qsa_block_topk wrote
|
||||
std::vector<int32_t> a(n_id), b(n_id);
|
||||
ck(cudaMemcpy(a.data(), ref_ids, n_id * 4, cudaMemcpyDeviceToHost));
|
||||
ck(cudaMemcpy(b.data(), two_ids, n_id * 4, cudaMemcpyDeviceToHost));
|
||||
int64_t bad_q = 0, compared = 0;
|
||||
for (int64_t i = 0; i < nq; ++i) {
|
||||
const int64_t n_kv = steps[(size_t) (i * k::kStepCount + k::kStepNKv)];
|
||||
const int64_t w = std::min<int64_t>(n_kv, steps[(size_t) (i * k::kStepCount + k::kStepWidth)]);
|
||||
compared += w;
|
||||
if (!std::equal(a.begin() + i * cap, a.begin() + i * cap + w, b.begin() + i * cap)) ++bad_q;
|
||||
}
|
||||
const bool ok = bad_q == 0 && st[tp::kChkIdQueries] == 0 && st[tp::kChkRescoredDiff] == 0 && st[tp::kChkKeptOut] == 0 &&
|
||||
st[tp::kChkQueries] == (unsigned long long) nq;
|
||||
float ratio;
|
||||
const uint32_t rb = (uint32_t) st[tp::kChkRatioMax];
|
||||
std::memcpy(&ratio, &rb, 4);
|
||||
const unsigned long long kept = st[tp::kStQueries] - st[tp::kStFull];
|
||||
std::printf("%s ctx=%7lld stride=%5lld queries=%3lld %-20s: ids %lld queries differ (%lld ids compared); scored "
|
||||
"again %llu (%.1f per query), %llu of %llu scored-again blocks not bitwise; flagged %llu (%llu tiles in "
|
||||
"full); kept %llu, outside the bound %llu, max |S1-S3|/eps %.4f\n",
|
||||
ok ? "ok " : "FAIL", (long long) c.ctx, (long long) max_blocks, (long long) nq, kKindName[c.kind],
|
||||
(long long) bad_q, (long long) compared, st[tp::kStBand], kept ? (double) st[tp::kStBand] / kept : 0.0,
|
||||
st[tp::kChkRescoredDiff], st[tp::kChkRescored], st[tp::kStFull], st[tp::kStFullTiles], st[tp::kChkKept],
|
||||
st[tp::kChkKeptOut], (double) ratio);
|
||||
if (ok && c.time && stride > 0) { // one synthetic observation of each scorer (the bench worker measures the engine)
|
||||
cudaEvent_t e0, e1, e2;
|
||||
ck(cudaEventCreate(&e0)); ck(cudaEventCreate(&e1)); ck(cudaEventCreate(&e2));
|
||||
constexpr int N = 5;
|
||||
ck(cudaEventRecord(e0));
|
||||
for (int r = 0; r < N; ++r) k::qsa_block_scores_tc(dp, dd, dq, dt, nq, max_blocks, s, ref, nullptr, active);
|
||||
ck(cudaEventRecord(e1));
|
||||
for (int r = 0; r < N; ++r)
|
||||
k::qsa_block_scores_two_pass(dp, dd, dq, dt, nq, nq_cap, max_blocks, s, two, scratch, sbytes, nullptr, active);
|
||||
ck(cudaEventRecord(e2));
|
||||
ck(cudaEventSynchronize(e2));
|
||||
float t3 = 0, t2 = 0;
|
||||
ck(cudaEventElapsedTime(&t3, e0, e1)); ck(cudaEventElapsedTime(&t2, e1, e2));
|
||||
std::printf(" time per call (synthetic, one run): 3xTF32 scorer %.3f ms, two-pass %.3f ms (keys once per chunk "
|
||||
"not included)\n", t3 / N, t2 / N);
|
||||
k::qsa_two_pass_stats(st, true, nullptr);
|
||||
cudaEventDestroy(e0); cudaEventDestroy(e1); cudaEventDestroy(e2);
|
||||
}
|
||||
ck(cudaFree(dp)); ck(cudaFree(dd)); ck(cudaFree(dq)); ck(cudaFree(dt)); ck(cudaFree(ref)); ck(cudaFree(two));
|
||||
ck(cudaFree(ref_ids)); ck(cudaFree(two_ids)); ck(cudaFree(scratch));
|
||||
return ok;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int main() {
|
||||
#if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906)
|
||||
// AMD: the two-pass scorer is CUDA-only (qsa_select.cu builds stubs under HIP), so the engine keeps today's scorer
|
||||
std::puts("SKIP: the two-pass scorer is CUDA-only");
|
||||
return 77;
|
||||
#endif
|
||||
#if defined(_WIN32)
|
||||
if (std::getenv("STRATA_SELECT_TWO_PASS") == nullptr) _putenv_s("STRATA_SELECT_TWO_PASS", "1");
|
||||
#else
|
||||
setenv("STRATA_SELECT_TWO_PASS", "1", 0); // read once, at the first call
|
||||
#endif
|
||||
int devices = 0;
|
||||
if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) {
|
||||
std::puts("SKIP: no CUDA device");
|
||||
return 77;
|
||||
}
|
||||
if (!k::qsa_two_pass_on()) {
|
||||
std::puts("FAIL: STRATA_SELECT_TWO_PASS is set to something other than 1");
|
||||
return 1;
|
||||
}
|
||||
const Case cases[] = {
|
||||
{2100, 64, kNormal, false}, {4096, 256, kNormal, false}, {9000, 37, kNormal, false},
|
||||
{20000, 256, kNormal, false}, {20481, 256, kSmall, false}, {32768, 256, kIndexer, false},
|
||||
{40000, 256, kWide, false}, {50000, 200, kNear, false}, {70000, 256, kDup, false},
|
||||
{131072, 256, kIndexer, false}, {131075, 129, kNormal, false}, {262144, 256, kIndexer, true},
|
||||
{262144, 256, kNormal, true}, {200003, 256, kSmall, false}, {65536, 256, kIndexer, true},
|
||||
};
|
||||
bool all = true;
|
||||
for (const int64_t stride : {(int64_t) 0, (int64_t) 65538})
|
||||
for (const Case& c : cases) all = run(c, stride) && all;
|
||||
if (!all) {
|
||||
std::puts("FAIL: the two-pass selection differs from the 3xTF32 scorer's");
|
||||
return 1;
|
||||
}
|
||||
std::puts("PASS: two-pass ids identical to the 3xTF32 scorer's, scored-again blocks bitwise, kept blocks inside the bound");
|
||||
return 0;
|
||||
}
|
||||
+108
-2
@@ -49,6 +49,8 @@
|
||||
#include "strata/kernels/qsa_pattn.hpp"
|
||||
#include "strata/kernels/qsa_prompt_attn.hpp"
|
||||
#include "strata/kernels/qsa_select.hpp"
|
||||
#include "strata/kernels/qsa_two_pass.hpp"
|
||||
#include "strata/kernels/qsa_two_pass_status.hpp"
|
||||
#include "strata/prefill/gemm.hpp"
|
||||
#include "strata/prefill/moe_fused.hpp"
|
||||
#include "strata/prefill/moe_ws.hpp"
|
||||
@@ -4382,6 +4384,22 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
|
||||
break;
|
||||
}
|
||||
pt.mark(kPfQsaSel, cs);
|
||||
#if !defined(STRATA_GPU_METAL)
|
||||
// STRATA_SELECT_TWO_PASS=1 (qsa_two_pass.hpp): the tensor-core scorer's calls whose last query has
|
||||
// STRATA_SELECT_TWO_PASS_MIN blocks or more score every block in single-pass TF32 and again in 3xTF32
|
||||
// only where the error bound reaches the cut-off; qsa_block_topk below returns the same cells. Its
|
||||
// scratch is m.attn, which the attention below writes and nothing reads before; the key norms are
|
||||
// this layer's and chunk's (the indexer appends above are complete)
|
||||
static const bool tp_on = strata::kernels::qsa_two_pass_on();
|
||||
static const bool tp_check = tp_on && strata::kernels::qsa_two_pass_check_on();
|
||||
static const int64_t tp_min = strata::kernels::qsa_two_pass_min_blocks();
|
||||
const uint64_t tp_bytes = (uint64_t) m.T * ZV * sizeof(float);
|
||||
const int64_t tp_nbid = T > 0 ? (int64_t) m.steps_host[(size_t) ((T - 1) * strata::kernels::kStepCount +
|
||||
strata::kernels::kStepNBid)] : 0;
|
||||
const bool tp_ok = tp_on && tp_nbid + 1 >= tp_min &&
|
||||
strata::kernels::qsa_two_pass_keys(st.idx_pooled, tp_nbid, m.sel_batch, m.max_blocks,
|
||||
m.attn, tp_bytes, m.cs);
|
||||
#endif
|
||||
for (int64_t t0 = 0; t0 < T; t0 += m.sel_batch) {
|
||||
const int64_t nb = std::min(m.sel_batch, T - t0);
|
||||
const int32_t* steps0 = m.steps_dev + t0 * strata::kernels::kStepCount;
|
||||
@@ -4391,13 +4409,66 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
|
||||
// the scores on tensor cores (3xTF32: FP32-level, not bitwise), below sm_80 the FP32 tiled kernel with STRATA_SELECT_SIMT=1;
|
||||
// STRATA_SELECT_OLD=1: the warp kernel
|
||||
static const bool old_sel = std::getenv("STRATA_SELECT_OLD") != nullptr;
|
||||
if (old_sel || !strata::kernels::qsa_block_scores_tc(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512,
|
||||
bool tp_call = false;
|
||||
#if !defined(STRATA_GPU_METAL)
|
||||
tp_call = tp_ok && !old_sel && active >= tp_min &&
|
||||
strata::kernels::qsa_block_scores_two_pass(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512,
|
||||
steps0, nb, m.sel_batch, m.max_blocks, s,
|
||||
m.sel_scores, m.attn, tp_bytes, m.cs, active);
|
||||
#endif
|
||||
if (!tp_call && (old_sel || !strata::kernels::qsa_block_scores_tc(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512,
|
||||
steps0, nb, m.max_blocks, s, m.sel_scores,
|
||||
m.cs, active))
|
||||
m.cs, active)))
|
||||
strata::kernels::qsa_block_scores(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512, steps0, nb,
|
||||
m.max_blocks, s, m.sel_scores, m.cs, active);
|
||||
strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s,
|
||||
m.sel_ids + t0 * m.cap, m.cs, active);
|
||||
#if !defined(STRATA_GPU_METAL)
|
||||
// STRATA_SELECT_TWO_PASS_CHECK=1 (diagnostic, one GPU): today's scores and ids for the same
|
||||
// queries, compared with the two-pass ones (counts at the prompt's end). Their buffers are the
|
||||
// rest of m.attn past the two-pass scratch when it is large enough (the attention below writes
|
||||
// m.attn after the comparison), else buffers of their own
|
||||
if (tp_call && tp_check) {
|
||||
static float* own_sc = nullptr;
|
||||
static int32_t* own_ids = nullptr;
|
||||
static size_t own_sc_n = 0, own_ids_n = 0;
|
||||
const size_t want_sc = (size_t) m.sel_batch * (size_t) m.max_blocks, want_ids = (size_t) m.sel_batch * (size_t) m.cap;
|
||||
const uint64_t at = (strata::kernels::qsa_two_pass_scratch_bytes(m.sel_batch, m.max_blocks) + 255u) & ~uint64_t(255u);
|
||||
const uint64_t at_ids = at + (((uint64_t) want_sc * 4u + 255u) & ~uint64_t(255u));
|
||||
float* ref_sc = nullptr;
|
||||
int32_t* ref_ids = nullptr;
|
||||
if (at_ids + (uint64_t) want_ids * 4u <= tp_bytes) {
|
||||
ref_sc = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(m.attn) + at);
|
||||
ref_ids = reinterpret_cast<int32_t*>(reinterpret_cast<uint8_t*>(m.attn) + at_ids);
|
||||
} else {
|
||||
if (want_sc > own_sc_n || want_ids > own_ids_n) {
|
||||
cudaStreamSynchronize(m.cs);
|
||||
cudaFree(own_sc);
|
||||
cudaFree(own_ids);
|
||||
own_sc = nullptr;
|
||||
own_ids = nullptr;
|
||||
own_sc_n = own_ids_n = 0;
|
||||
if (cudaMalloc((void**) &own_sc, want_sc * 4) != cudaSuccess ||
|
||||
cudaMalloc((void**) &own_ids, want_ids * 4) != cudaSuccess) {
|
||||
err = "STRATA_SELECT_TWO_PASS_CHECK: no room for its buffers";
|
||||
return false;
|
||||
}
|
||||
own_sc_n = want_sc;
|
||||
own_ids_n = want_ids;
|
||||
}
|
||||
ref_sc = own_sc;
|
||||
ref_ids = own_ids;
|
||||
}
|
||||
if (!strata::kernels::qsa_block_scores_tc(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512, steps0, nb,
|
||||
m.max_blocks, s, ref_sc, m.cs, active)) {
|
||||
err = "STRATA_SELECT_TWO_PASS_CHECK: the tensor-core scorer did not run";
|
||||
return false;
|
||||
}
|
||||
strata::kernels::qsa_block_topk(ref_sc, steps0, nb, m.max_blocks, m.cap, s, ref_ids, m.cs, active);
|
||||
strata::kernels::qsa_two_pass_check(m.sel_scores, m.sel_ids + t0 * m.cap, ref_sc, ref_ids, steps0, nb,
|
||||
m.sel_batch, m.max_blocks, m.cap, m.attn, m.cs);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
// STRATA_SEL_OVERLAP (debug, D-1's question): how much do neighbouring queries' selections share?
|
||||
// Per tile of 16 queries: the union of their selected cells against the sum of their widths.
|
||||
@@ -6298,6 +6369,31 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
|
||||
(long long) kv_prefetches);
|
||||
}
|
||||
stats_.ms_total += ms_since(t_start);
|
||||
#if !defined(STRATA_GPU_METAL)
|
||||
if (strata::kernels::qsa_two_pass_on()) { // STRATA_SELECT_TWO_PASS: what the two-pass scorer did in this prompt
|
||||
namespace tp = strata::kernels::two_pass;
|
||||
unsigned long long c[tp::kStatCount] = {};
|
||||
if (strata::kernels::qsa_two_pass_stats(c, true, m.cs) && c[tp::kStCalls] > 0) {
|
||||
const unsigned long long sparse = c[tp::kStQueries], kept = sparse - c[tp::kStFull];
|
||||
std::fprintf(stderr, "strata prefill: two-pass QSA select: %llu calls, %llu queries (%llu blocks) in single-pass TF32, "
|
||||
"%llu blocks again in 3xTF32 (%.1f per query, max %llu), %llu queries flagged (%llu "
|
||||
"16-query tiles in full), %llu selecting every cell; candidates %.0f per query (max %llu)\n",
|
||||
c[tp::kStCalls], sparse, c[tp::kStBlocks], c[tp::kStBand], kept ? (double) c[tp::kStBand] / kept : 0.0,
|
||||
c[tp::kStBandMax], c[tp::kStFull], c[tp::kStFullTiles], c[tp::kStIdentity],
|
||||
kept ? (double) c[tp::kStCand] / kept : 0.0, c[tp::kStCandMax]);
|
||||
if (c[tp::kChkQueries] > 0) {
|
||||
const uint32_t rb = (uint32_t) c[tp::kChkRatioMax];
|
||||
float ratio = 0.0f;
|
||||
std::memcpy(&ratio, &rb, 4);
|
||||
std::fprintf(stderr, "strata prefill: two-pass check: %llu queries, %llu with ids not the 3xTF32 scorer's (%llu ids); "
|
||||
"scored again: %llu of %llu blocks not the 3xTF32 score bit for bit; kept: %llu of %llu "
|
||||
"single-pass scores outside the bound, the largest |S1 - S3| / bound %.4g\n",
|
||||
c[tp::kChkQueries], c[tp::kChkIdQueries], c[tp::kChkIdEntries], c[tp::kChkRescoredDiff],
|
||||
c[tp::kChkRescored], c[tp::kChkKeptOut], c[tp::kChkKept], (double) ratio);
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
{
|
||||
static const bool f16_range = std::getenv("STRATA_F16_RANGE") != nullptr;
|
||||
float xm = 0.0f, wy[2] = {0.0f, 0.0f}; unsigned long long xo = 0, wyo[2] = {0, 0};
|
||||
@@ -7800,7 +7896,12 @@ bool Prefill::reserve_uploads(int64_t T) {
|
||||
}
|
||||
|
||||
bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) {
|
||||
#if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906)
|
||||
(void) strata::kernels::qsa_two_pass_on(); // HIP: a set STRATA_SELECT_TWO_PASS says once it is not used
|
||||
#endif
|
||||
#if defined(STRATA_GPU_METAL)
|
||||
static strata::kernels::two_pass::RequestStatus tp_status(std::getenv("STRATA_SELECT_TWO_PASS"));
|
||||
(void) tp_status.on("Metal");
|
||||
{
|
||||
const auto t0 = std::chrono::steady_clock::now();
|
||||
// DFlash2: the window writes the taps of the prompt rows at positions >= tap_first_ and hands them to on_taps after each chunk
|
||||
@@ -7839,6 +7940,11 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
|
||||
}
|
||||
}
|
||||
if (lm_prepare_chain(n, why) > 0) {
|
||||
#if !defined(STRATA_USE_HIP) && !defined(STRATA_HIP_GFX906)
|
||||
// STRATA_SELECT_TWO_PASS is in run_impl's selection only, not in this walk: say so once
|
||||
static strata::kernels::two_pass::RequestStatus tp_status(std::getenv("STRATA_SELECT_TWO_PASS"));
|
||||
(void) tp_status.on("the layer-major prompt path");
|
||||
#endif
|
||||
bool ok = run_layer_major(tokens, n, pos0, err);
|
||||
std::string drain_err; // a layer split: the later stages' last super-chunk
|
||||
if (!drain_pipeline(drain_err) && ok) {
|
||||
|
||||
@@ -37,6 +37,7 @@
|
||||
#include "strata/kernels/qsa_decode_attn.hpp"
|
||||
#include "strata/kernels/qsa_prompt_attn.hpp"
|
||||
#include "strata/kernels/qsa_select.hpp"
|
||||
#include "strata/kernels/qsa_two_pass_status.hpp"
|
||||
#include "strata/prefill/gemm.hpp"
|
||||
#include "strata/prefill/moe_fused.hpp"
|
||||
#include "strata/prefill/moe_fused_iq.hpp"
|
||||
@@ -5196,6 +5197,8 @@ bool Prefill::drain_pipeline(std::string& err) {
|
||||
}
|
||||
|
||||
bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) {
|
||||
static strata::kernels::two_pass::RequestStatus tp_status(std::getenv("STRATA_SELECT_TWO_PASS"));
|
||||
(void) tp_status.on("SYCL"); // the two-pass scorer is CUDA-only: say once that a set switch is not used
|
||||
const bool body_ok = run_impl(tokens, n, pos0, err);
|
||||
|
||||
std::string drain_err;
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
// STRATA_SELECT_TWO_PASS's request and its line where the two-pass scorer cannot run (qsa_two_pass_status.hpp).
|
||||
// GPU-free: the parse, quiet when unset or supported, one line per process where unsupported.
|
||||
#include "strata/kernels/qsa_two_pass_status.hpp"
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <string>
|
||||
|
||||
namespace tp = strata::kernels::two_pass;
|
||||
|
||||
static void check(bool ok, const char* why) {
|
||||
if (!ok) {
|
||||
std::fprintf(stderr, "FAIL: %s\n", why);
|
||||
std::exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
static std::string output(FILE* log) {
|
||||
check(std::fflush(log) == 0 && std::fseek(log, 0, SEEK_SET) == 0, "rewind the log");
|
||||
std::string text;
|
||||
char buf[512];
|
||||
while (const size_t n = std::fread(buf, 1, sizeof(buf), log)) text.append(buf, n);
|
||||
check(!std::ferror(log), "read the log");
|
||||
return text;
|
||||
}
|
||||
|
||||
static std::string line(const char* where) {
|
||||
return std::string("strata: STRATA_SELECT_TWO_PASS=1 is not used on ") + where +
|
||||
": the two-pass scorer runs only in the CUDA prompt path's tensor-core selection (sm_80+); the selection is "
|
||||
"unchanged\n";
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
check(argc == 2, "pass a scratch log path (no GPU or model needed)");
|
||||
FILE* log = std::fopen(argv[1], "w+");
|
||||
check(log != nullptr, "open the log");
|
||||
|
||||
// exactly "1" is on, as the scorer parsed it before; unset or anything else is off and prints nothing anywhere
|
||||
check(tp::requested_from("1"), "1: on");
|
||||
const char* off[] = {nullptr, "", "0", "01", "1x", "true", "1 ", "2"};
|
||||
for (const char* flag : off) {
|
||||
check(!tp::requested_from(flag), "not exactly 1: off");
|
||||
tp::RequestStatus status(flag);
|
||||
check(!status.requested(), "not requested");
|
||||
check(!status.on("HIP", log), "off stays off where unsupported");
|
||||
check(!status.on(nullptr, log), "off stays off where supported");
|
||||
}
|
||||
check(output(log).empty(), "no line without the switch");
|
||||
|
||||
// a supported path returns the request and prints nothing (the engine's per-segment line reports what ran)
|
||||
tp::RequestStatus cuda("1");
|
||||
check(cuda.on(nullptr, log) && cuda.on(nullptr, log), "supported: on");
|
||||
check(output(log).empty(), "supported: no line");
|
||||
|
||||
// each unsupported path the engine names: off, one line, never repeated
|
||||
const char* where[] = {"HIP", "Metal", "SYCL", "this CUDA card or with STRATA_QSA_WARP",
|
||||
"the layer-major prompt path"};
|
||||
std::string expected;
|
||||
for (const char* w : where) {
|
||||
tp::RequestStatus status("1");
|
||||
check(!status.on(w, log), "unsupported: off");
|
||||
expected += line(w);
|
||||
check(output(log) == expected, "unsupported: one line naming the path");
|
||||
check(!status.on(w, log) && !status.on(w, log), "unsupported again: off");
|
||||
check(output(log) == expected, "unsupported: the line is printed once");
|
||||
}
|
||||
|
||||
// a layer split can mix cards: a supported query neither uses up the line nor is turned off by an unsupported one
|
||||
tp::RequestStatus mixed("1");
|
||||
check(mixed.on(nullptr, log), "supported card first");
|
||||
check(!mixed.on("this CUDA card or with STRATA_QSA_WARP", log), "then an unsupported card");
|
||||
check(mixed.on(nullptr, log), "the supported card stays on");
|
||||
expected += line("this CUDA card or with STRATA_QSA_WARP");
|
||||
check(output(log) == expected, "the unsupported card's line once");
|
||||
|
||||
check(std::fclose(log) == 0, "close the log");
|
||||
check(std::remove(argv[1]) == 0, "remove the log");
|
||||
std::puts("qsa_two_pass_status_test: OK");
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
// tests/prefill/qsa_two_pass_test.cpp - STRATA_SELECT_TWO_PASS without a GPU (qsa_two_pass.hpp).
|
||||
//
|
||||
// 1. TF32: the host copy of cvt.rna.tf32.f32 rounds to nearest, ties away from zero.
|
||||
// 2. The bound: for random and adversarial indexer queries and keys, the single-pass score S1 and the 3xTF32 score S3
|
||||
// are computed from the same TF32 operands as the kernels, with each MMA modelled two ways - exact then rounded to
|
||||
// nearest (the most accurate hardware), and aligned to the largest addend with every addend and the result
|
||||
// truncated to 24 bits (no guard bits, the least accurate model) - and |S1 - S3| must stay within eps. The exact
|
||||
// cross term |qh.kl + ql.kh| must stay within its Cauchy-Schwarz part.
|
||||
// 3. The cut-off rule: rows of 3xTF32 scores, single-pass scores anywhere inside [S3 - eps, S3 + eps] (pushed across
|
||||
// the threshold where possible), ties at the threshold, every tail weight, short and long rows (the sampled
|
||||
// candidate threshold), the identity: a host copy of the band kernel picks the blocks to score again, and the
|
||||
// top-k of the resulting row (block_topk_kernel's rule) must be the top-k of the 3xTF32 row, id for id.
|
||||
// Exits non-zero on the first failure. No GPU, no model.
|
||||
#include "strata/kernels/qsa_two_pass.hpp"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
namespace tp = strata::kernels::two_pass;
|
||||
|
||||
namespace {
|
||||
|
||||
int g_fail = 0;
|
||||
#define CHECK(cond, ...) \
|
||||
do { \
|
||||
if (!(cond)) { \
|
||||
std::fprintf(stderr, "FAIL %s:%d: ", __FILE__, __LINE__); \
|
||||
std::fprintf(stderr, __VA_ARGS__); \
|
||||
std::fprintf(stderr, "\n"); \
|
||||
++g_fail; \
|
||||
if (g_fail > 20) std::exit(1); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
float up(double x) { // the smallest float >= x (x finite and in range)
|
||||
float f = (float) x;
|
||||
if ((double) f < x) f = std::nextafter(f, std::numeric_limits<float>::infinity());
|
||||
return f;
|
||||
}
|
||||
float down(double x) {
|
||||
float f = (float) x;
|
||||
if ((double) f > x) f = std::nextafter(f, -std::numeric_limits<float>::infinity());
|
||||
return f;
|
||||
}
|
||||
|
||||
// Shewchuk's exact summation (Python's math.fsum): the sum of doubles, correctly rounded
|
||||
double fsum(const std::vector<double>& xs) {
|
||||
std::vector<double> p;
|
||||
for (double x : xs) {
|
||||
size_t i = 0;
|
||||
for (double y : p) {
|
||||
if (std::fabs(x) < std::fabs(y)) std::swap(x, y);
|
||||
const double hi = x + y, lo = y - (hi - x);
|
||||
if (lo != 0.0) p[i++] = lo;
|
||||
x = hi;
|
||||
}
|
||||
p.resize(i);
|
||||
p.push_back(x);
|
||||
}
|
||||
double s = 0.0;
|
||||
for (size_t i = p.size(); i-- > 0;) s += p[i]; // the partials are non-overlapping: this is the rounded sum
|
||||
return s;
|
||||
}
|
||||
|
||||
// ---- 1. TF32
|
||||
void test_tf32() {
|
||||
auto bits = [](float f) { uint32_t b; std::memcpy(&b, &f, 4); return b; };
|
||||
CHECK(tp::tf32_rna_host(1.0f) == 1.0f, "1.0");
|
||||
CHECK(tp::tf32_rna_host(1.0f + 0x1p-11f) == 1.0f + 0x1p-10f, "a tie rounds away from zero");
|
||||
CHECK(tp::tf32_rna_host(-(1.0f + 0x1p-11f)) == -(1.0f + 0x1p-10f), "a negative tie rounds away from zero");
|
||||
CHECK(tp::tf32_rna_host(1.0f + 0x1p-12f) == 1.0f, "below the tie rounds down");
|
||||
CHECK((bits(tp::tf32_rna_host(3.14159265f)) & 0x1fffu) == 0u, "the low 13 bits are zero");
|
||||
std::mt19937 rng(7);
|
||||
std::normal_distribution<float> nd;
|
||||
for (int i = 0; i < 100000; ++i) {
|
||||
const float x = nd(rng) * std::ldexp(1.0f, (int) (rng() % 40) - 20);
|
||||
const float h = tp::tf32_rna_host(x);
|
||||
CHECK(std::fabs((double) x - (double) h) <= std::ldexp(std::fabs((double) x), -11), "TF32 error %g", (double) x);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 2. the bound
|
||||
struct Op { float h, l; }; // tf32(x), tf32(x - tf32(x)): the kernels' split
|
||||
Op split(float x) {
|
||||
const float h = tp::tf32_rna_host(x);
|
||||
return {h, tp::tf32_rna_host(x - h)};
|
||||
}
|
||||
|
||||
// one MMA: the accumulator c plus 8 exact products
|
||||
float mma_rn(float c, const double* p) {
|
||||
std::vector<double> t(p, p + 8);
|
||||
t.push_back(c);
|
||||
return (float) fsum(t);
|
||||
}
|
||||
float trunc_to(double v, int e) { // toward zero to a multiple of 2^e
|
||||
const double q = std::ldexp(1.0, e);
|
||||
return (float) (std::trunc(v / q) * q);
|
||||
}
|
||||
float mma_trunc(float c, const double* p) {
|
||||
double t[9];
|
||||
for (int i = 0; i < 8; ++i) t[i] = p[i];
|
||||
t[8] = c;
|
||||
int emax = std::numeric_limits<int>::min();
|
||||
for (double v : t)
|
||||
if (v != 0.0) { int e; std::frexp(v, &e); emax = std::max(emax, e); }
|
||||
if (emax == std::numeric_limits<int>::min()) return 0.0f;
|
||||
std::vector<double> a;
|
||||
for (double v : t) { // every addend kept to 24 bits below the largest one's leading bit, toward zero
|
||||
const double q = std::ldexp(1.0, emax - 24);
|
||||
a.push_back(std::trunc(v / q) * q);
|
||||
}
|
||||
const double s = fsum(a);
|
||||
if (s == 0.0) return 0.0f;
|
||||
int e;
|
||||
std::frexp(s, &e);
|
||||
return trunc_to(s, e - 24); // the result truncated to FP32's 24 bits
|
||||
}
|
||||
|
||||
struct Pair { std::vector<float> q, k; }; // q: 4 heads x 128, k: 128
|
||||
|
||||
// The scores the kernels compute, per model: S1 (qh*kh, 16 MMAs per head) and S3 (ql*kh, qh*kl, qh*kh per step).
|
||||
void scores(const Pair& x, bool trunc, float& s1, float& s3, double cross[4], double hn[4][2], double kn[2]) {
|
||||
std::vector<Op> k(128);
|
||||
for (int i = 0; i < 128; ++i) k[i] = split(x.k[i]);
|
||||
double kh2 = 0, kl2 = 0;
|
||||
for (int i = 0; i < 128; ++i) { kh2 += (double) k[i].h * k[i].h; kl2 += (double) k[i].l * k[i].l; }
|
||||
kn[0] = std::sqrt(kh2);
|
||||
kn[1] = std::sqrt(kl2);
|
||||
s1 = s3 = 0.0f;
|
||||
for (int h = 0; h < 4; ++h) {
|
||||
std::vector<Op> q(128);
|
||||
for (int i = 0; i < 128; ++i) q[i] = split(x.q[h * 128 + i]);
|
||||
double qh2 = 0, ql2 = 0;
|
||||
for (int i = 0; i < 128; ++i) { qh2 += (double) q[i].h * q[i].h; ql2 += (double) q[i].l * q[i].l; }
|
||||
hn[h][0] = std::sqrt(qh2);
|
||||
hn[h][1] = std::sqrt(ql2);
|
||||
std::vector<double> d;
|
||||
for (int i = 0; i < 128; ++i) { d.push_back((double) q[i].h * k[i].l); d.push_back((double) q[i].l * k[i].h); }
|
||||
cross[h] = fsum(d);
|
||||
float c1 = 0.0f, c3 = 0.0f;
|
||||
for (int st = 0; st < 16; ++st) {
|
||||
double p_hh[8], p_lh[8], p_hl[8];
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
const int i = 8 * st + j;
|
||||
p_hh[j] = (double) q[i].h * k[i].h;
|
||||
p_lh[j] = (double) q[i].l * k[i].h;
|
||||
p_hl[j] = (double) q[i].h * k[i].l;
|
||||
}
|
||||
c1 = trunc ? mma_trunc(c1, p_hh) : mma_rn(c1, p_hh);
|
||||
c3 = trunc ? mma_trunc(c3, p_lh) : mma_rn(c3, p_lh); // the kernel's order: al*bh, ah*bl, ah*bh
|
||||
c3 = trunc ? mma_trunc(c3, p_hl) : mma_rn(c3, p_hl);
|
||||
c3 = trunc ? mma_trunc(c3, p_hh) : mma_rn(c3, p_hh);
|
||||
}
|
||||
s1 += c1 > 0.0f ? c1 : 0.0f;
|
||||
s3 += c3 > 0.0f ? c3 : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// eps as the band kernel forms it (norms rounded up, the coefficients and eps rounded up)
|
||||
float eps_of(const double hn[4][2], const double kn[2]) {
|
||||
double a = 0, c = 0;
|
||||
for (int h = 0; h < 4; ++h) { a += hn[h][0]; c += hn[h][1]; }
|
||||
a = up(a) * (1 + 1e-12);
|
||||
c = up(c) * (1 + 1e-12);
|
||||
double P, Q;
|
||||
tp::coefficients(a, c, P, Q);
|
||||
const float Pf = up(P), Qf = up(Q), kh = up(kn[0] * (1 + 1e-12)), kl = up(kn[1] * (1 + 1e-12));
|
||||
return up((double) Pf * kl + (double) Qf * kh + (double) tp::kAbs);
|
||||
}
|
||||
|
||||
void test_bound() {
|
||||
std::mt19937 rng(20261024);
|
||||
std::normal_distribution<float> nd;
|
||||
std::uniform_int_distribution<int> small(-3, 3), expo(-12, 12), coin(0, 9);
|
||||
// the worst rounding: mantissas just below a TF32 tie, so |x - tf32(x)| is nearly 2^-11 |x|
|
||||
auto worst = [&](float sign) {
|
||||
uint32_t b = 0x3f800000u | ((rng() & 0x3ffu) << 13) | 0x0fffu; // ...0 1111 1111 1111 in the dropped bits
|
||||
float f;
|
||||
std::memcpy(&f, &b, 4);
|
||||
return sign * std::ldexp(f, expo(rng) / 4);
|
||||
};
|
||||
const char* names[] = {"normal", "wide magnitudes and zeros", "small integers", "aligned worst rounding (q = k)",
|
||||
"worst rounding, signs of k"};
|
||||
for (int mode = 0; mode < 5; ++mode) {
|
||||
double worst_ratio[2] = {0, 0}, worst_cross = 0;
|
||||
const int N = mode >= 3 ? 1500 : 3000;
|
||||
for (int n = 0; n < N; ++n) {
|
||||
Pair x;
|
||||
x.q.resize(512);
|
||||
x.k.resize(128);
|
||||
for (int i = 0; i < 128; ++i) {
|
||||
float v;
|
||||
if (mode == 0) v = nd(rng);
|
||||
else if (mode == 1) v = coin(rng) == 0 ? 0.0f : std::ldexp(nd(rng), expo(rng));
|
||||
else if (mode == 2) v = (float) small(rng);
|
||||
else v = worst(nd(rng) > 0 ? 1.0f : -1.0f);
|
||||
x.k[i] = v;
|
||||
}
|
||||
for (int h = 0; h < 4; ++h)
|
||||
for (int i = 0; i < 128; ++i) {
|
||||
float v;
|
||||
if (mode == 0) v = nd(rng);
|
||||
else if (mode == 1) v = coin(rng) == 0 ? 0.0f : std::ldexp(nd(rng), expo(rng));
|
||||
else if (mode == 2) v = (float) small(rng);
|
||||
else if (mode == 3) v = x.k[i]; // every head aligned with k, lo parts too: Cauchy-Schwarz is tight
|
||||
else v = worst(x.k[i] > 0 ? 1.0f : -1.0f); // independent worst roundings, positive dot products
|
||||
x.q[h * 128 + i] = v;
|
||||
}
|
||||
for (int model = 0; model < 2; ++model) {
|
||||
float s1, s3;
|
||||
double cross[4], hn[4][2], kn[2];
|
||||
scores(x, model == 1, s1, s3, cross, hn, kn);
|
||||
const float e = eps_of(hn, kn);
|
||||
const double d = std::fabs((double) s1 - (double) s3);
|
||||
CHECK(d <= (double) e, "%s, %s model: |S1 - S3| = %.9g > eps %.9g", names[mode],
|
||||
model ? "truncating" : "rounding", d, (double) e);
|
||||
worst_ratio[model] = std::max(worst_ratio[model], d / (double) e);
|
||||
if (model == 0)
|
||||
for (int h = 0; h < 4; ++h) {
|
||||
const double cb = hn[h][0] * kn[1] + hn[h][1] * kn[0];
|
||||
CHECK(std::fabs(cross[h]) <= cb * (1 + 1e-12) + 1e-300, "%s: cross term %.9g > %.9g", names[mode],
|
||||
std::fabs(cross[h]), cb);
|
||||
if (cb > 0) worst_cross = std::max(worst_cross, std::fabs(cross[h]) / cb);
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("bound, %-36s: largest |S1 - S3| / eps %.4f (rounding MMA), %.4f (truncating MMA); "
|
||||
"cross term / its bound %.4f\n", names[mode], worst_ratio[0], worst_ratio[1], worst_cross);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 3. the cut-off rule
|
||||
constexpr int R = 4;
|
||||
uint32_t key(float s) { return tp::order_key_host(s); }
|
||||
|
||||
// the largest key k with (sum of weights of entries with key >= k) >= width (block_topk_kernel's threshold)
|
||||
uint32_t kth(std::vector<std::pair<uint32_t, int>> kw, int64_t width) {
|
||||
std::sort(kw.begin(), kw.end(), [](auto& a, auto& b) { return a.first > b.first; });
|
||||
int64_t c = 0;
|
||||
for (auto& [k, w] : kw) {
|
||||
c += w;
|
||||
if (c >= width) return k;
|
||||
}
|
||||
return 0u; // fewer cells than the width: the radix scan ends at digit 0 every pass
|
||||
}
|
||||
|
||||
// block_topk_kernel's ids: cells of blocks above the threshold, and the first (width - above) tied cells by index
|
||||
std::vector<int32_t> topk(const std::vector<float>& sc, int64_t n_kv, int64_t n_bid, int64_t width) {
|
||||
std::vector<int32_t> out;
|
||||
if (n_kv <= width) {
|
||||
for (int64_t j = 0; j < n_kv; ++j) out.push_back((int32_t) j);
|
||||
return out;
|
||||
}
|
||||
auto weight = [&](int64_t b) { return b < n_bid ? R : (int) (n_kv - n_bid * R); };
|
||||
std::vector<std::pair<uint32_t, int>> kw;
|
||||
for (int64_t b = 0; b <= n_bid; ++b)
|
||||
if (weight(b)) kw.push_back({key(sc[b]), weight(b)});
|
||||
const uint32_t thr = kth(kw, width);
|
||||
int64_t above = 0;
|
||||
for (int64_t b = 0; b <= n_bid; ++b)
|
||||
if (weight(b) && key(sc[b]) > thr) above += weight(b);
|
||||
int64_t eq_left = width - above;
|
||||
for (int64_t b = 0; b <= n_bid; ++b) {
|
||||
const int w = weight(b);
|
||||
if (!w) continue;
|
||||
const uint32_t k = key(sc[b]);
|
||||
if (k > thr) for (int c = 0; c < w; ++c) out.push_back((int32_t) (b * R + c));
|
||||
else if (k == thr) for (int c = 0; c < w && eq_left > 0; ++c, --eq_left) out.push_back((int32_t) (b * R + c));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
struct Band { int mode = 0; std::vector<int64_t> again; int64_t cand = 0; };
|
||||
|
||||
// the band kernel on the host (tp_band_kernel, qsa_select.cu): the same steps and constants
|
||||
Band band(const std::vector<float>& s1, const std::vector<float>& eps, int64_t n_kv, int64_t n_bid, int64_t width) {
|
||||
Band r;
|
||||
if (n_kv <= width) return r;
|
||||
const int64_t nb = n_bid + 1;
|
||||
const int w_tail = (int) (n_kv - n_bid * R);
|
||||
auto interval = [&](int64_t b, float& lo, float& hi, float& e) {
|
||||
e = eps[b];
|
||||
lo = down((double) s1[b] - (double) e);
|
||||
hi = up((double) s1[b] + (double) e);
|
||||
};
|
||||
float F = -std::numeric_limits<float>::infinity();
|
||||
if (nb > tp::kCandCap) {
|
||||
int64_t stride = tp::kSampleStride;
|
||||
while ((nb + stride - 1) / stride > tp::kCandCap) stride *= 2;
|
||||
const int64_t ns = (nb + stride - 1) / stride;
|
||||
std::vector<std::pair<uint32_t, int>> kw;
|
||||
for (int64_t i = 0; i < ns; ++i) {
|
||||
const int64_t b = i * stride;
|
||||
uint32_t k = 0u;
|
||||
if (b < n_bid) { float lo, hi, e; interval(b, lo, hi, e); k = key(lo); }
|
||||
else if (w_tail > 0) k = key(s1[b]);
|
||||
kw.push_back({k, 1});
|
||||
}
|
||||
const int64_t ks = 2 * ((width + R * stride - 1) / (R * stride)) + 16;
|
||||
const uint32_t fk = kth(kw, std::min(ks, ns));
|
||||
if (fk != 0u) F = tp::key_float_host(fk);
|
||||
}
|
||||
std::vector<int64_t> cand;
|
||||
int64_t wf = 0;
|
||||
float E = 0.0f;
|
||||
for (int64_t b = 0; b < nb; ++b) {
|
||||
float lo, hi, e = 0.0f;
|
||||
int w;
|
||||
if (b < n_bid) { w = R; interval(b, lo, hi, e); }
|
||||
else { w = w_tail; if (!w) continue; lo = hi = s1[b]; }
|
||||
if (lo >= F) wf += w;
|
||||
if (hi >= F) { cand.push_back(b); E = std::max(E, e); }
|
||||
}
|
||||
r.cand = (int64_t) cand.size();
|
||||
if ((int64_t) cand.size() > tp::kCandCap || wf < width) { r.mode = 2; return r; }
|
||||
std::vector<std::pair<uint32_t, int>> kw;
|
||||
for (int64_t b : cand) kw.push_back({key(s1[b]), b < n_bid ? R : w_tail});
|
||||
const float T1 = tp::key_float_host(kth(kw, width));
|
||||
const float lo_b = std::max(down((double) T1 - (double) E), F), hi_b = up((double) T1 + (double) E);
|
||||
int res[8] = {};
|
||||
for (int64_t b : cand) {
|
||||
if (b >= n_bid) continue;
|
||||
float lo, hi, e;
|
||||
interval(b, lo, hi, e);
|
||||
if (hi >= lo_b && lo <= hi_b) { r.again.push_back(b); ++res[b & 7]; }
|
||||
}
|
||||
for (int x : res) if (x > tp::kResCap) { r.mode = 2; return r; }
|
||||
r.mode = 1;
|
||||
return r;
|
||||
}
|
||||
|
||||
void test_rule() {
|
||||
std::mt19937 rng(424242);
|
||||
std::normal_distribution<double> nd;
|
||||
std::uniform_real_distribution<double> ud(0.0, 1.0);
|
||||
const int64_t sizes[] = {300, 512, 600, 1500, 4000, 5119, 5120, 5121, 9000, 20000, 65000};
|
||||
const char* kinds[] = {"continuous", "small integers (ties)", "a plateau at the threshold", "eps wide against the spread",
|
||||
"heavy tail", "every score equal"};
|
||||
int64_t rows = 0, full = 0, again = 0, cand = 0;
|
||||
for (int kind = 0; kind < 6; ++kind) {
|
||||
const int64_t rows0 = rows, full0 = full, again0 = again;
|
||||
for (int64_t n_bid : sizes) {
|
||||
for (int tail = 0; tail < 4; ++tail) {
|
||||
for (int rep = 0; rep < (n_bid > 20000 ? 2 : 4); ++rep) {
|
||||
const int64_t n_kv = n_bid * R + tail, width = std::min<int64_t>(n_kv, 2048 + R - 1);
|
||||
std::vector<float> s3((size_t) n_bid + 1), s1((size_t) n_bid + 1), eps((size_t) n_bid + 1);
|
||||
// kind 0: continuous; 1: small integers (ties); 2: a plateau at the threshold; 3: eps wide
|
||||
// against the spread; 4: heavy tail; 5: every score equal
|
||||
for (int64_t b = 0; b < n_bid; ++b) {
|
||||
double v;
|
||||
if (kind == 1) v = (double) (int) (ud(rng) * 40);
|
||||
else if (kind == 2) v = ud(rng) < 0.3 ? 7.0 : 7.0 + nd(rng);
|
||||
else if (kind == 4) v = std::exp(2.0 * nd(rng));
|
||||
else if (kind == 5) v = 3.0;
|
||||
else v = std::max(0.0, 2.0 + nd(rng));
|
||||
s3[b] = (float) v;
|
||||
const double scale = kind == 3 ? 0.3 : (kind == 1 ? 0.6 : 0.002);
|
||||
eps[b] = (float) (scale * (0.5 + ud(rng)));
|
||||
}
|
||||
s3[n_bid] = tail ? 1e9f + (float) ud(rng) : (float) ud(rng); // the tail: exact, +1e9 with cells
|
||||
eps[n_bid] = 0.0f;
|
||||
// the 3xTF32 threshold, to push single-pass scores across it
|
||||
const std::vector<int32_t> ref = topk(s3, n_kv, n_bid, width);
|
||||
std::vector<std::pair<uint32_t, int>> kw;
|
||||
for (int64_t b = 0; b <= n_bid; ++b) {
|
||||
const int w = b < n_bid ? R : tail;
|
||||
if (w) kw.push_back({key(s3[b]), w});
|
||||
}
|
||||
const float T = n_kv > width ? tp::key_float_host(kth(kw, width)) : 0.0f;
|
||||
for (int64_t b = 0; b < n_bid; ++b) {
|
||||
// S1 inside [S3 - eps, S3 + eps]: towards and across the threshold where it can, else random
|
||||
const double e = eps[b] * (1.0 - 1e-6);
|
||||
double d;
|
||||
if (rep % 2 == 0) d = (s3[b] >= T ? -e : e);
|
||||
else d = (2.0 * ud(rng) - 1.0) * e;
|
||||
s1[b] = (float) ((double) s3[b] + d);
|
||||
if (std::fabs((double) s1[b] - (double) s3[b]) > eps[b]) s1[b] = s3[b];
|
||||
}
|
||||
s1[n_bid] = s3[n_bid];
|
||||
const Band bd = band(s1, eps, n_kv, n_bid, width);
|
||||
std::vector<float> row = s1;
|
||||
if (bd.mode == 2) row = s3;
|
||||
for (int64_t b : bd.again) row[b] = s3[b];
|
||||
const std::vector<int32_t> got = topk(row, n_kv, n_bid, width);
|
||||
CHECK(got == ref, "kind %d, %lld blocks, tail %d, rep %d: ids differ (mode %d, %zu scored again)", kind,
|
||||
(long long) n_bid, tail, rep, bd.mode, bd.again.size());
|
||||
++rows;
|
||||
full += bd.mode == 2;
|
||||
again += (int64_t) bd.again.size();
|
||||
cand += bd.cand;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("rule, %-28s: %3lld rows, %3lld scored again in full, %.1f blocks scored again per row\n", kinds[kind],
|
||||
(long long) (rows - rows0), (long long) (full - full0),
|
||||
rows > rows0 ? (double) (again - again0) / (double) (rows - rows0) : 0.0);
|
||||
}
|
||||
std::printf("rule: %lld rows, ids identical in every one; %lld scored again in full (the list or the bound could not "
|
||||
"settle them), %.1f blocks scored again and %.0f candidates per row\n", (long long) rows, (long long) full,
|
||||
rows ? (double) again / rows : 0.0, rows ? (double) cand / rows : 0.0);
|
||||
// the guard bites: keep the single-pass score of one block whose interval reaches the threshold, and the ids move
|
||||
{
|
||||
const int64_t n_bid = 1000, n_kv = n_bid * R, width = 2048 + R - 1;
|
||||
std::vector<float> s3((size_t) n_bid + 1), s1((size_t) n_bid + 1), eps((size_t) n_bid + 1, 0.01f);
|
||||
for (int64_t b = 0; b < n_bid; ++b) s3[b] = 1.0f + 0.001f * (float) b; // distinct, rising
|
||||
s3[n_bid] = 0.0f;
|
||||
eps[n_bid] = 0.0f;
|
||||
const std::vector<int32_t> ref = topk(s3, n_kv, n_bid, width);
|
||||
// the 513th block from the top (block 487) sits at the threshold; its single-pass score 0.009 lower
|
||||
for (int64_t b = 0; b <= n_bid; ++b) s1[b] = s3[b];
|
||||
s1[487] = s3[487] - 0.009f;
|
||||
s1[486] = s3[486] + 0.009f;
|
||||
const Band bd = band(s1, eps, n_kv, n_bid, width);
|
||||
CHECK(bd.mode == 1 && std::find(bd.again.begin(), bd.again.end(), 487) != bd.again.end(),
|
||||
"block 487 must be scored again (mode %d)", bd.mode);
|
||||
std::vector<float> row = s1;
|
||||
for (int64_t b : bd.again) row[b] = s3[b];
|
||||
CHECK(topk(row, n_kv, n_bid, width) == ref, "the guarded row must select today's ids");
|
||||
CHECK(topk(s1, n_kv, n_bid, width) != ref, "without the re-score the ids must differ (the test bites)");
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main() {
|
||||
test_tf32();
|
||||
test_bound();
|
||||
test_rule();
|
||||
if (g_fail) {
|
||||
std::fprintf(stderr, "FAILED: %d checks\n", g_fail);
|
||||
return 1;
|
||||
}
|
||||
std::puts("PASS: qsa_two_pass_test (TF32 split, the bound under two MMA models, the cut-off rule)");
|
||||
return 0;
|
||||
}
|
||||
Reference in new issue
Block a user