Foundry/tt/kernels/resident_topk_protocol.h

503 lines
24 KiB
C

// SPDX-License-Identifier: AGPL-3.0-or-later
//
// resident_topk_protocol.h
// The device half of ft_tt_op_sample.h: the parameter block, the result block, the
// per-turn sampler state, the ORDER COMPARATOR, and the pseudo-random generator, shared
// by the resident TOPK/SAMPLE program's three images. The host half is ft_tt_op_sample.h;
// the two must agree word for word and tests/test_topk_ops.c pins the pair off-hardware
// by PARSING this file rather than by retyping it.
//
// ===========================================================================
// 1. WHY THIS OP EXISTS AT ALL, AND WHY IT CANNOT BE A HOST FUNCTION
// ===========================================================================
// Card -> host is 2.3 MB/s at 8-byte loads and 4.7 MB/s at 16 (dev notes §6), and
// device -> host DMA does not exist on Blackhole. MiniCPM5-1B's vocabulary is
// 130,560 (MEASURED, read out of the production artifact's `output.weight` directory
// entry -- not the 73k this op was scoped against), so ONE bf16 logit row is 261,120
// bytes and costs 55-113 ms to fetch. Decode cannot pay that per token, and the 2-core
// no-AVX host cannot afford the 130k-wide softmax on the other side of it either.
//
// So the requirement is a shape, not a preference: **the only thing that crosses the link
// per token is a 4-byte token id.** Everything above -- selection, the penalty, the
// temperature, the two nucleus filters and the draw -- happens here.
//
// ===========================================================================
// 2. THE SHAPE: A STREAMING SELECTION NETWORK, NOT A TILE MAP
// ===========================================================================
// Every other op in this backend is reader -> compute -> writer with the TRISCs doing the
// arithmetic. Top-k is not that op and should not be forced into it, for three
// independent reasons:
//
// a. IT RETURNS AN INDEX. The FPU and SFPU produce VALUES; `reduce_tile<MAX>` gives the
// row maximum and cannot say which column it came from. Blackhole does ship a bitonic
// top-k SFPU (`topk_local_sort`/`topk_merge`/`topk_rebuild`, plus
// `max_reduce_with_indices`) which carries index tiles alongside the data -- but it
// sorts 64 values spanning TWO tiles along the column axis, wants K in {4,8,16,32,64},
// runs in APPROXIMATE mode, and needs the sorted run to live in DST. Our row lies
// along W, is 4080 tiles wide, and Wt is not a power of two.
// b. THE ROW DOES NOT FIT L1. 4080 tiles x 2048 B is 8,355,840 B against a 1,351,680 B
// CB arena (ft_tt_op_l1_layout.h). Any exact answer is therefore a STREAMING
// reduction with bounded state, whatever unit computes it -- 6.2x too big is not a
// tuning problem.
// c. THE STATE IS TINY AND SCALAR. N <= 64 (index, key) pairs is 512 bytes. A sorted
// insert is a branch and a shift, which is what a RISC-V is for and what a 32-lane
// SIMD unit is not.
//
// Hence: **the reader block-prefetches logit tiles into a circular buffer and the
// data-movement processor beside it runs a scalar streaming selection over them.** The
// compute thread is deliberately idle (resident_topk_compute.c says so in its own
// header, and says what would put it back to work).
//
// ===========================================================================
// 3. THE HOT LOOP IS ONE SUBTRACT AND ONE UNSIGNED COMPARE PER LOGIT
// ===========================================================================
// The naive scalar test is "map this bf16 to an order-preserving integer, compare against
// the current N-th best". The map (§4) is four operations, and at 130,560 elements a token
// that is the whole budget.
//
// It is not needed on the REJECT path, which is 99.97% of elements. "Does this bf16 beat
// the incumbent `worst`" is a test on the RAW BITS, and the winning set is always ONE
// CONTIGUOUS INTERVAL of the 16-bit pattern space:
//
// worst is non-negative (bits < 0x8000): a candidate wins iff it is positive AND its
// bits are larger, i.e. bits in (worst, 0x8000).
// worst is negative (bits >= 0x8000): every positive wins, and a negative wins iff its
// bits are SMALLER (larger bits = larger magnitude = more negative), i.e. the
// winning set is bits in [0, worst) -- the two ranges are contiguous.
//
// So the whole test is `(uint32_t)(bits - accept_low) < accept_span` with the pair
// recomputed only when `worst` changes -- ~ONE subtract and ONE branch per logit. The
// order-preserving map is then paid only by the winners. Measured effect is in the report;
// the point here is that the reject path never converts a float.
//
// ===========================================================================
// 4. THE ORDER, AND WHY TIES ARE THE MAIN EVENT AND NOT AN EDGE CASE
// ===========================================================================
// bf16 has EIGHT mantissa bits. Over a 130,560-element row, exact ties are certain, not
// exotic -- there are only 65,536 bf16 patterns, so the pigeonhole gives an average of two
// logits per representable value before you look at the distribution. Tie handling is
// therefore the dominant correctness question for this op, which is why the acceptance
// gate is exact index agreement and not a probability score: a top-k whose values are all
// right and whose indices are off by a tie correlates at 1.0 with the truth and is wrong.
//
// THE COMPARATOR, matching Pantheon's K80 `fc_topk_logits` (topk_logits.cu:46-50) and the
// host serial scan it reproduces (`if (logits[i] > logits[best]) best = i`):
//
// a beats b <=> key(a) > key(b) || (key(a) == key(b) && index(a) < index(b))
//
// i.e. **LOWEST INDEX WINS A TIE**. Scanning in ascending index order, that is exactly
// "insert only on strictly greater", which is what ftk_select_offer() does.
//
// key() is the standard order-preserving map from IEEE bits to unsigned:
// key(u) = (u & 0x8000) ? (uint16)~u : (uint16)(u | 0x8000)
//
// KNOWN AND DELIBERATE: key() orders -0.0 (0x8000 -> 0x7FFF) BELOW +0.0 (0x0000 -> 0x8000),
// where IEEE says they are equal. The host reference is defined on the same key order so
// the gate is exact, and tests/test_topk_ops.c sweeps ALL 65,536 bf16 patterns to prove
// key order and float order agree everywhere else. NaN is refused, not ordered: the host
// checks the row and this kernel does not, because a per-element NaN test is exactly the
// four operations §3 exists to avoid.
//
// ===========================================================================
// 5. WHAT A CLOSED DECODE LOOP HAS TO CARRY BETWEEN TOKENS
// ===========================================================================
// This is the reason the sampler's knobs live in a STATE block on the card and not in the
// descriptor. A descriptor field is per-call; every one of these is per-TURN or
// per-SEQUENCE, and a host that had to refresh them per token would be back in the path
// the op exists to leave. ftk_state is that list, and it is deliberately the whole list:
//
// - the PRNG state, so successive tokens are not the same draw and so a turn is
// reproducible from a seed the host set once;
// - the sampler knobs (temperature, top_k, top_p, repeat penalty and its window), which
// change when a request changes, not when a token is emitted;
// - the RECENT-TOKEN RING and its cursor, appended to BY THE SAMPLER, because the
// repetition penalty's input is the loop's own output;
// - the EMITTED-TOKEN RING and its count, so the host drains N tokens per poll instead
// of being in the path of each one;
// - the stop condition -- an EOS id set and a token budget -- because a loop the host
// must stop is a loop the host is in.
//
// The sampled id is ALSO written to `token_out`, an address the descriptor names, which is
// where the next step's embedding lookup reads its input. That is the actual closure: the
// token never leaves the card to come back as the next input.
//
#ifndef FOUNDRY_RESIDENT_TOPK_PROTOCOL_H
#define FOUNDRY_RESIDENT_TOPK_PROTOCOL_H
#include <stdint.h>
/* Header-only helper linkage, spelled the same way in both languages -- `inline` in C++ and
* `static inline` in C. NOT cosmetic: internal linkage changes GCC's inlining and cloning
* decisions, and the COMPUTE images that also include this header must not move for a
* change that is about the data-movement ones. */
#ifndef FTQ_INLINE
#ifdef __cplusplus
#define FTQ_INLINE inline
#else
#define FTQ_INLINE static inline
#endif
#endif
#ifndef __cplusplus
#include <stdbool.h>
#endif
// ---------------------------------------------------------------------------
// The sixteen parameter words. Mirrors ft_tt_op_sample.h -- change both or neither.
// ---------------------------------------------------------------------------
/** Tiles in the logit row (Wt). 4080 for MiniCPM5-1B: 130,560 / 32, exactly. */
enum { kTopkTileCount = 0 };
/** Which of a tile's 32 rows carries the token being decoded, 0..31. The LM head runs at
* Mt = 1, so 32 token rows are produced and one is wanted; this names it. */
enum { kTopkRow = 1 };
/** How many (index, value) pairs to select. <= kTopkMaxSelect. */
enum { kTopkSelectCount = 2 };
/** Tiles per NoC barrier in the reader; must DIVIDE kTopkTileCount or a block straddles
* the CB wrap. The kernel repeats the host's arithmetic rather than trusting it. */
enum { kTopkDepth = 3 };
/** The logit row in DRAM: a NoC endpoint from ft_tt_coord, never a guess (spec §7.3). */
enum { kTopkLogitsX = 4 };
enum { kTopkLogitsY = 5 };
enum { kTopkLogitsBase = 6 };
/** The result block, in THIS core's L1. TOPK writes it, SAMPLE reads it. */
enum { kTopkResultAddress = 7 };
/** The per-turn sampler state, in THIS core's L1. See §5. */
enum { kTopkStateAddress = 8 };
/** Where SAMPLE writes the chosen token id -- normally the next step's embedding input, so
* the loop closes on the card. A zero address means "result block only". It names a
* kTopkTokenSlotBytes SLOT, not a word: see the hand-off note below. */
enum { kTopkTokenOutX = 9 };
enum { kTopkTokenOutY = 10 };
enum { kTopkTokenOutBase = 11 };
/**
* Measurement flags. A RUNTIME word and not a build flag, on purpose: the A/B arms of the
* cost split have to be the SAME BINARY or the measurement is of two implementations
* rather than of the thing being measured (docs/tt_kernel_model.md §14, the
* `weight_n_major` precedent).
*/
enum { kTopkFlags = 12 };
/** Pop every logit block WITHOUT selecting over it. The arm that measures the DRAM stream
* underneath the scan: full-arm minus this arm IS the scalar selection's cost. It writes
* an empty result block, so nothing can mistake it for an answer. */
enum { kTopkFlagDrainOnly = 1 };
/** Reserved; written 0 by the host and never read, exactly as the vector family's spare
* words are. NOT free for a future op to redefine. */
enum { kTopkWords = 16 };
/** Most pairs one result block holds. 64 keeps the block inside one 512-byte L1 region and
* covers every top_k Pantheon serves (its production value is 40). */
enum { kTopkMaxSelect = 64 };
// ---------------------------------------------------------------------------
// The result block: a header the host can refuse, then N pairs
// ---------------------------------------------------------------------------
// Descending by value, ties by ascending index -- the ORDER `fc_topk_logits` returns, so a
// host that already consumes the K80's output consumes this unchanged. SAMPLE re-sorts to
// ascending index internally, because Pantheon's sparse sampler scans the vocabulary in
// index order and the cumulative draw is order-sensitive (engine.c sparse_sample()).
enum { kTopkResultMagic = 0 }; // kTopkResultMagicValue; refuses a wrong address
enum { kTopkResultCount = 1 }; // pairs actually written
enum { kTopkResultSequence = 2 }; // ++ per TOPK, so a stale block is visible
enum { kTopkResultTokenValid = 3 }; // SAMPLE writes 1 here after a draw
enum { kTopkResultToken = 4 }; // the sampled id, mirrored for the host
enum { kTopkResultHeaderWords = 8 };
/** Two words per pair: [vocabulary index, value as fp32 BITS]. The value is stored as fp32
* rather than bf16 because bf16 -> fp32 is exact (a 16-bit left shift) and every consumer
* -- this kernel's own sampler and the host reference both -- does fp32 arithmetic. */
enum { kTopkPairWords = 2 };
enum { kTopkPairIndex = 0 };
enum { kTopkPairValueBits = 1 };
enum { kTopkResultMagicValue = 0x46544B31 }; // "FTK1"
// ---------------------------------------------------------------------------
// THE TOKEN HAND-OFF IS AN ALIGNED 16-BYTE NoC WRITE, AND THAT COST TWO CORES
// ---------------------------------------------------------------------------
// The obvious way to hand four bytes to the next op is `noc_inline_dw_write`, and it is
// wrong here for a documented reason: **"The destination node can
// be either a Tensix core+L1 memory address or a PCIe controller; this API does not support
// DRAM addresses"** (the stock dataflow API; its watcher build has a sanitizer whose
// only job is to catch this). `token_out` normally names
// DRAM -- the embedding table's index array is not in anyone's L1 -- so the op was using the
// one destination class the call excludes.
//
// It does not fail loudly. On Blackhole a non-register destination takes the SPOOF path
// (the stock inline-write routine): the value is staged in
// a 16-byte L1 scratch slot and issued as a **4-byte** `noc_async_write`. The token LANDS --
// the on-card check that reads it back passed every time -- and the core is left unable to
// boot: the next `assert_reset` + firmware load comes up with BRISC out of reset, hanging
// before it releases its subordinates (`soft reset 0x00047000`, `ft_tt_core status`'s
// "boot handshake never finished"). MEASURED on Tensix 20 and 21, bisected to this one write
// on Tensix 27, which survives the identical run with the write suppressed.
//
// Four bytes is also below what this backend's OWN transfer contract allows: a DRAM WRITE
// takes the L1 alignment, 16 B in address AND in length (`ft_tt_noc_alignment()`,
// ft_tt_noc.c). `ft_tt_noc_check_transfer()` would have refused this transfer -- but it runs
// on the HOST, and nothing re-checks on the kernel side, which is where the transfer is
// actually issued. So the hand-off is now a legal transfer instead of a forbidden call:
// **16 bytes, 16-byte aligned at both ends, through plain `noc_async_write`** -- legal to L1
// and to DRAM alike, and needing no knowledge of which one the descriptor named (the kernel
// has an (x, y) and cannot tell).
//
// The consequence for a caller: `token_out` names a 16-byte SLOT. The id lands in its first
// word and the other three are zeroed.
/** First word of the 16-byte source slot the hand-off writes FROM, inside the result block.
* It IS kTopkResultToken, so the mirrored id and the shipped id are the same word and
* cannot disagree. */
enum { kTopkTokenSlotWords = kTopkResultToken };
/** Bytes moved by the hand-off, and the alignment both ends must satisfy. */
enum { kTopkTokenSlotBytes = 16 };
enum { kTopkTokenOutAlignment = 16 };
FTQ_STATIC_ASSERT((kTopkTokenSlotWords * 4u) % kTopkTokenOutAlignment == 0u,
"the hand-off's L1 source must be aligned within a result block that is");
FTQ_STATIC_ASSERT(kTopkTokenSlotWords + kTopkTokenSlotBytes / 4u <= kTopkResultHeaderWords,
"the hand-off slot must stay inside the result header, clear of the pairs");
/** Bytes of a result block that holds kTopkMaxSelect pairs. */
enum { kTopkResultBytes = (kTopkResultHeaderWords + kTopkMaxSelect * kTopkPairWords) * 4u };
// ---------------------------------------------------------------------------
// The sampler state block (§5)
// ---------------------------------------------------------------------------
enum { kSampleStateMagic = 0 };
/** xorshift32. Never zero -- the host refuses a zero seed, because xorshift32 has a fixed
* point there and a "random" sampler that returns the same draw forever looks like a
* temperature bug for a long time. */
enum { kSampleStateRngState = 1 };
/** All four knobs are fp32 BITS, so the block is a plain uint32 array on both sides and
* the host never has to agree with the device about float endianness or promotion. */
enum { kSampleStateTemperatureBits = 2 };
enum { kSampleStateTopPBits = 3 };
enum { kSampleStateRepeatPenaltyBits = 4 };
enum { kSampleStateTopK = 5 };
enum { kSampleStateRepeatWindow = 6 }; // entries of the recent ring to consult
enum { kSampleStateRecentCursor = 7 }; // next write position in the recent ring
enum { kSampleStateRecentCount = 8 }; // entries written so far, saturating
enum { kSampleStateEmittedCount = 9 }; // tokens this state has produced
enum { kSampleStateEosCount = 10 };
enum { kSampleStateStopped = 11 }; // 1 once EOS or the budget was hit
enum { kSampleStateTokenBudget = 12 };
enum { kSampleStateLastToken = 13 };
enum { kSampleStateHeaderWords = 16 };
enum { kSampleStateMagicValue = 0x46544B53 }; // "FTKS"
/** The recent-token ring, appended to by the sampler itself. */
enum { kSampleRecentOffsetWords = kSampleStateHeaderWords };
enum { kSampleRecentCapacity = 64 };
/** The EOS id set. Small and linear; a model has one or two. */
enum { kSampleEosOffsetWords = kSampleRecentOffsetWords + kSampleRecentCapacity };
enum { kSampleEosCapacity = 4 };
/** The emitted-token ring the host drains lazily. */
enum { kSampleEmittedOffsetWords = kSampleEosOffsetWords + kSampleEosCapacity };
enum { kSampleEmittedCapacity = 256 };
enum { kSampleStateWords = kSampleEmittedOffsetWords + kSampleEmittedCapacity };
enum { kSampleStateBytes = kSampleStateWords * 4u };
// ---------------------------------------------------------------------------
// The order (§4)
// ---------------------------------------------------------------------------
/** bf16 bits -> an unsigned key that orders exactly as the float does. Branchless: the
* sign-broadcast form, so the winners' path has no unpredictable branch either. */
FTQ_INLINE uint32_t ftk_key_from_bf16(uint32_t bits)
{
const uint32_t sign_mask = (uint32_t)((int32_t)(bits << 16) >> 31);
return (bits ^ 0x8000u ^ (sign_mask & 0x7FFFu)) & 0xFFFFu;
}
/** bf16 bits -> fp32 bits. Exact: bf16 IS the top half of fp32. */
FTQ_INLINE uint32_t ftk_fp32_bits_from_bf16(uint32_t bits)
{
return (bits & 0xFFFFu) << 16;
}
/**
* The reject-path interval of §3: the set of raw bf16 patterns that BEAT `worst_bits`,
* as one contiguous unsigned range `[low, low + span)`.
*
* `full` asks for the range that accepts everything, which is what the selection uses
* before it has N entries.
*/
FTQ_INLINE void ftk_accept_range(uint32_t worst_bits, bool accept_all, uint32_t *low,
uint32_t *span)
{
if (accept_all) {
*low = 0u;
*span = 0x10000u;
} else if (worst_bits < 0x8000u) {
*low = worst_bits + 1u;
*span = 0x8000u - *low;
} else {
*low = 0u;
*span = worst_bits;
}
}
/** One selected pair, held in the selector's registers/stack rather than in L1. */
typedef struct ftk_entry {
uint32_t key;
uint32_t index;
uint32_t bf16_bits;
} ftk_entry;
/**
* The selection state: `entries` sorted DESCENDING by key, ties by ascending index.
*
* `accept_low` / `accept_span` are the §3 interval, kept in step with `entries[count-1]`
* and recomputed only when the N-th best changes.
*/
typedef struct ftk_select {
ftk_entry entries[kTopkMaxSelect];
uint32_t count;
uint32_t capacity;
uint32_t accept_low;
uint32_t accept_span;
} ftk_select;
FTQ_INLINE void ftk_select_init(ftk_select *state, uint32_t capacity)
{
state->count = 0u;
state->capacity = (capacity > (uint32_t)kTopkMaxSelect) ? (uint32_t)kTopkMaxSelect : capacity;
ftk_accept_range(0u, true, &state->accept_low, &state->accept_span);
}
FTQ_INLINE void ftk_select_refresh_range(ftk_select *state)
{
const bool accept_all = state->count < state->capacity;
ftk_accept_range(accept_all ? 0u : state->entries[state->capacity - 1u].bf16_bits, accept_all,
&state->accept_low, &state->accept_span);
}
/**
* Offer one logit to the selection. THE ONLY place the comparator lives.
*
* Insertion is "walk back while the incumbent is strictly smaller", so an EQUAL key stops
* the walk and the candidate lands AFTER it. Scanning in ascending index order, that is
* exactly lowest-index-wins (§4) -- the tie rule is a property of where the walk stops,
* not a separate test, which is why it cannot be forgotten in one of two branches.
*/
FTQ_INLINE void ftk_select_offer(ftk_select *state, uint32_t bf16_bits, uint32_t index)
{
const uint32_t key = ftk_key_from_bf16(bf16_bits);
uint32_t position = (state->count < state->capacity) ? state->count : state->capacity - 1u;
if (state->count < state->capacity)
state->count++;
while (position > 0u && state->entries[position - 1u].key < key) {
state->entries[position] = state->entries[position - 1u];
position--;
}
state->entries[position].key = key;
state->entries[position].index = index;
state->entries[position].bf16_bits = bf16_bits;
ftk_select_refresh_range(state);
}
/** Whether `bf16_bits` is worth offering. One subtract and one unsigned compare (§3). */
FTQ_INLINE bool ftk_select_admits(const ftk_select *state, uint32_t bf16_bits)
{
return (bf16_bits - state->accept_low) < state->accept_span;
}
// ---------------------------------------------------------------------------
// The generator
// ---------------------------------------------------------------------------
/** xorshift32 (Marsaglia 2003, the 13/17/5 triple). Chosen because the device and the host
* reference must produce the SAME stream for the sampled-token gate to be exact, and a
* libc `rand()` cannot be that on a kernel with no libc. */
FTQ_INLINE uint32_t ftk_rng_next(uint32_t *state)
{
uint32_t value = *state;
value ^= value << 13;
value ^= value >> 17;
value ^= value << 5;
*state = value;
return value;
}
/** A draw in [0, 1), matching the reference's `(float)rand() / (RAND_MAX + 1.0f)` in RANGE
* and in half-openness. 24 bits, because that is fp32's mantissa and anything finer is
* discarded by the conversion anyway. */
FTQ_INLINE float ftk_rng_next_unit_float(uint32_t *state)
{
return (float)(ftk_rng_next(state) >> 8) * (1.0f / 16777216.0f);
}
// ---------------------------------------------------------------------------
// exp, in software, because the data-movement RISC-V has no FPU
// ---------------------------------------------------------------------------
// MEASURED: the data-movement -mcpu is `tt-bh` = `rv32im_zmmul_zaamo_zba_zbb`, ilp32 --
// integer only, no F extension. So every float operation here is a libgcc call and there is
// no libm to reach for. `expf` is therefore not available and is not merely slow.
//
// exp(x) = 2^(x*log2 e), split into an integer power of two applied to the exponent field
// and a degree-6 minimax polynomial for 2^f on f in [-0.5, 0.5]. The coefficients are
// PARSED out of this file by tests/test_topk_ops.c and compared against the host
// reference's, which is what keeps the two implementations from drifting -- a retyped
// constant is exactly the failure mode dev notes §4 catalogues.
//
// Only x <= 0 is ever passed (the sampler always subtracts the row maximum first), so the
// overflow half of the range is a refusal rather than a saturation.
static const float kFtkLog2E = 1.44269504088896340736f;
static const float kFtkExp2C1 = 0.69314718055994531f;
static const float kFtkExp2C2 = 0.24022650695910072f;
static const float kFtkExp2C3 = 0.05550410866482158f;
static const float kFtkExp2C4 = 0.00961812910762848f;
static const float kFtkExp2C5 = 0.00133335581467035f;
static const float kFtkExp2C6 = 0.00015403530393381f;
/** exp(x) for x <= 0. Returns 0.0f on underflow, exactly as expf() does. */
FTQ_INLINE float ftk_exp_nonpositive(float x)
{
float scaled = x * kFtkLog2E;
int32_t power;
float fraction;
float poly;
uint32_t exponent_bits;
float scale;
if (scaled < -126.0f)
return 0.0f;
// `scaled` is non-positive, so truncation-toward-zero of (scaled - 0.5) IS the nearest
// integer -- no rounding-mode call, and |fraction| <= 0.5 by construction.
power = (int32_t)(scaled - 0.5f);
fraction = scaled - (float)power;
poly = kFtkExp2C6;
poly = poly * fraction + kFtkExp2C5;
poly = poly * fraction + kFtkExp2C4;
poly = poly * fraction + kFtkExp2C3;
poly = poly * fraction + kFtkExp2C2;
poly = poly * fraction + kFtkExp2C1;
poly = poly * fraction + 1.0f;
exponent_bits = (uint32_t)(power + 127) << 23;
__builtin_memcpy(&scale, &exponent_bits, sizeof(scale));
return poly * scale;
}
#endif // FOUNDRY_RESIDENT_TOPK_PROTOCOL_H