503 lines
24 KiB
C
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
|