Add a Head¶
A head consumes a conditioning vector and returns an action. Follow the flow head.
The contract¶
class MyHead {
MyHead(std::span<const TensorView> weights, Arena& scratch) noexcept;
bool valid() const noexcept;
void sample(std::span<const float> cond, // conditioning vector, size d_model
std::span<const float> noise, // action_dim
std::size_t steps, Method m,
std::span<float> out) noexcept; // action chunk
};
The constructor resolves weight pointers by name and derives config from their shapes. sample reads the conditioning vector and writes the action. It allocates nothing.
Steps¶
Create
src/core/heads/<name>/<name>.{h,cc}. Resolve weight pointers in the constructor. Set anok_flag only when every required tensor is present.Add any missing kernel to
kernels.h. If it is public, implement it inkernels_avx2.cc,kernels_neon.cc, andkernels_scalar.cc, or the fallback build will not link.Carve scratch with
arena_span. Takemark()at the top ofsampleandreset_toat the end. Nothing allocates on the hot path.Add a converter mapping so a checkpoint loads. See Converter. Carry the normalization stats and un-normalize the action if the checkpoint stores them.
Add a sampling entrypoint to the C-ABI, or a head selector on
fe_engine_sample.Add a PyTorch reference to
tools/benchmark/torch_ref.pyand wire the ULP gate.Add tests to
test/core_model_tests.cc(ortest/kernel_tests.ccfor primitives). Add an example. Tick the head on the home page cards and write an ADR.
Reference: src/core/heads/flow/flow.h and flow.cc.