Spaces:
Sleeping
Sleeping
| namespace cism { | |
| using DotKernel = float (*)(const float*, const float*, std::size_t); | |
| using Int8DotKernel = float (*)(const std::int8_t*, const float*, std::size_t); | |
| // IEEE-754 binary16 conversion (round-to-nearest-even), portable software | |
| // implementation for packing and scalar fallback (no F16C requirement). | |
| std::uint16_t fp32_to_fp16(float value); | |
| float fp16_to_fp32(std::uint16_t bits); | |
| using Fp16DotKernel = float (*)(const std::uint16_t*, const float*, std::size_t); | |
| // int4/fp4 dots consume a permuted activation stream (even/odd 32-blocks, | |
| // built by the caller) for full blocks plus the original row for the scalar | |
| // tail; the 5th/4th scale args are per-32 FP32 / per-16 E4M3 bytes. | |
| using Int4DotKernel = float (*)(const std::uint8_t*, const float*, const float*, const float*, std::size_t); | |
| // FP4 (E2M1 elements, E4M3 scale per 16 elements): nibbles, permuted + original | |
| // activations, and scale bytes. | |
| using Fp4DotKernel = float (*)(const std::uint8_t*, const float*, const float*, const std::uint8_t*, std::size_t); | |
| // Quantized-activation kernels: the caller quantizes the shared activation | |
| // to int16 per 32-element block (32767/absmax, block-dequant scales passed | |
| // alongside); weights decode to int16 through pshufb LUTs and multiply | |
| // accumulate in integer pmaddwd lanes with independent accumulators. | |
| using Int8Q8DotKernel = float (*)(const std::int8_t*, const std::int16_t*, const float*, std::size_t); | |
| using Int4Q8DotKernel = float (*)(const std::uint8_t*, const std::int16_t*, const float*, const float*, std::size_t); | |
| using Fp4Q8DotKernel = float (*)(const std::uint8_t*, const std::int16_t*, const std::uint8_t*, const float*, std::size_t); | |
| DotKernel fp32_kernel(); | |
| Int8DotKernel int8_kernel(); | |
| Fp16DotKernel fp16_kernel(); | |
| Int4DotKernel int4_kernel(); | |
| Fp4DotKernel fp4_kernel(); | |
| // Quantized-activation selectors: null when AVX2 is unavailable (callers | |
| // fall back to the fp32-activation kernels above). | |
| Int8Q8DotKernel int8_q8_kernel(); | |
| Int4Q8DotKernel int4_q8_kernel(); | |
| Fp4Q8DotKernel fp4_q8_kernel(); | |
| const char* kernel_name(); | |
| // Portable-ISA dispatch (Intel/AMD/x86 + ARM phones). | |
| // Highest SIMD tier this binary can actually execute on this CPU: | |
| // "scalar" (any CPU) < "avx2" (x86-64, FMA) < "avx-vnni" (Zen4+/ADL+, | |
| // integer MAC fast path) < "avx512-vnni" (server AVX512) / "neon" (ARM64 | |
| // phones, Apple Silicon, Raspberry Pi). kernel_name() above keeps reporting | |
| // the fp32-kernel family ("scalar"/"avx2"/"neon"); kernel_variant() reports | |
| // this overall tier for telemetry and tests. | |
| const char* kernel_variant(); | |
| bool has_avx2_cpu(); | |
| bool has_f16c_cpu(); | |
| bool has_sse41_cpu(); | |
| bool has_avx_cpu(); | |
| bool has_avx_vnni_cpu(); | |
| bool has_avx512_vnni_cpu(); | |
| bool has_neon_cpu(); | |
| // int8-activation VNNI path: quantize_row_i8 shares quantize_row's blocking | |
| // (int8 127/absmax per 32, one fp32 dequant scale per block) so the VNNI | |
| // kernel sees the same blocking the int16 path uses. Selector is null unless | |
| // the binary was built with VNNI support AND the CPU has it (otherwise the | |
| // int16 Q8 path above is used, so pre-VNNI machines are untouched). | |
| void quantize_row_i8(const float* input, std::size_t n, std::int8_t* values, float* scales); | |
| void quantize_row_i8_avx2(const float* input, std::size_t n, std::int8_t* values, float* scales); | |
| using Int8Q8VnniKernel = float (*)(const std::int8_t*, const std::int8_t*, const float*, std::size_t); | |
| Int8Q8VnniKernel int8_q8_vnni_kernel(); | |
| float dot_scalar(const float* a, const float* b, std::size_t n); | |
| float dot_fp16_scalar(const std::uint16_t* weights, const float* input, std::size_t n); | |
| float dot_int8_scalar(const std::int8_t* weights, const float* input, std::size_t n); | |
| float dot_int4_scalar(const std::uint8_t* weights, const float* input, const float* scales, std::size_t n); | |
| float dot_fp4_scalar(const std::uint8_t* weights, const float* input, const std::uint8_t* scales, std::size_t n); | |
| // 5-arg adapters over the 4-arg scalars (ignore the permuted buffer) so the | |
| // permuted kernel type keeps a working non-AVX2 fallback. | |
| float dot_int4_scalar5(const std::uint8_t* weights, const float* act_perm, | |
| const float* input, const float* scales, std::size_t n); | |
| float dot_fp4_scalar5(const std::uint8_t* weights, const float* act_perm, | |
| const float* input, const std::uint8_t* scales, std::size_t n); | |
| // E2M1 quantization support shared by the packer and the kernels. | |
| // element_lut: raw nibble -> E2M1 value times two as int8 (exact small ints). | |
| // scale_lut: E4M3 bit pattern -> float scale, already halved so a decoded | |
| // element (half value) times the table entry equals E2M1 * E4M3 exactly. | |
| // encode_scale: positive float -> E4M3 bit pattern, round-half-even like | |
| // torch.float8_e4m3fn, saturating at 448. | |
| const std::int8_t* fp4_element_lut(); | |
| // Kernel tuning notes (Zen3 AVX2+FMA, generic /arch:AVX2, no VNNI): | |
| // - int8 is bandwidth-bound; 4-bit is kernel-bound (~26 uops/32 structural). | |
| // - Inner blocking below (64-wide unroll) and T0 prefetches are the only | |
| // kernel-internal levers that keep scalar-fallback bitwise semantics. | |
| // - Row-slice / thread-threshold / GEMM-blocking / SiLU-fusion live in | |
| // runtime.cpp (not owned here); this header documents the shared constants | |
| // so a future runtime wiring stays consistent. SIMD order may differ from | |
| // scalar; numeric changes stay PPL-gated (see VALIDATION.md noise note). | |
| constexpr std::size_t kParallelThreshold = 1u << 16; // rows*cols, runtime-owned | |
| constexpr std::size_t kDotUnroll = 64; // kernel inner blocking | |
| constexpr std::size_t kPrefetchDistance = 512; // bytes, T0, never faults | |
| // Permuted-activation helper for int4/fp4 (even/odd 32-blocks). Scalar | |
| // semantics; the AVX2 build offers an unrolled variant. Not wired to the | |
| // runtime path here (session-persistent scratch already reuses the buffer); | |
| // provided so future runtime wiring has one canonical implementation. | |
| void permute_act32_blocked(const float* input, std::size_t n, float* out); | |
| // Gated SiLU*up epilogue (stable sigmoid, matches Session::forward_tokens). | |
| // Test-gated only: silu_fusion_enabled() is false unless CISM_FUSE_SILU=1, | |
| // and the runtime path does not call these (prior A/B showed fusion ~7% | |
| // slower on Zen3 streaming hardware). Scalar is the bitwise reference. | |
| void silu_mul_scalar(const float* gate, const float* up, float* out, std::size_t n); | |
| bool silu_fusion_enabled(); | |
| // Vector-exp activation blocks for the decode hot path (Surjo GDN/MLP/XSA). | |
| // AVX2 TU evaluates exp with a ≤1-ULP polynomial (deterministic: same input | |
| // bits always give same output bits); the scalar fallback uses libm and is | |
| // the bitwise reference on non-AVX2 machines. In-place forms allow x==out. | |
| // Unified stable sigmoid 1/(1+exp(-x)) (exact for x>=0, ≤1 ULP elsewhere). | |
| // Numeric changes vs libm stay PPL-gated (VALIDATION.md noise note). | |
| void act_exp(float* x, std::size_t n); | |
| void act_silu(float* x, std::size_t n); | |
| void act_silu_mul(const float* gate, const float* up, float* out, std::size_t n); | |
| void act_sigmoid(float* x, std::size_t n); | |
| void act_sigmoid_mul(float* o, const float* g, std::size_t n); | |
| // Clamp-free SiLU*up for the dense path (no +-15 clamp there). | |
| void act_silu_mul_plain(const float* gate, const float* up, float* out, std::size_t n); | |
| // Exact GELU for the FWKV path (0.5*x*(1+erf(x/sqrt(2))), PyTorch exact). | |
| void act_gelu(float* x, std::size_t n); | |
| void permute_act32_avx2(const float* input, std::size_t n, float* out); | |
| void silu_mul_avx2(const float* gate, const float* up, float* out, std::size_t n); | |
| const float* fp4_scale_lut(); | |
| float fp4_decode_scale(std::uint8_t bits); | |
| std::uint8_t fp4_encode_scale(float value); | |
| float dot_avx2(const float* a, const float* b, std::size_t n); | |
| float dot_fp16_avx2(const std::uint16_t* weights, const float* input, std::size_t n); | |
| // Vector-exp act blocks (defined in kernels_avx2.cpp; scalar twins below). | |
| void act_exp_avx2(float* x, std::size_t n); | |
| void act_silu_avx2(float* x, std::size_t n); | |
| void act_silu_mul_avx2(const float* gate, const float* up, float* out, std::size_t n); | |
| void act_silu_mul_plain_avx2(const float* gate, const float* up, float* out, std::size_t n); | |
| void act_gelu_avx2(float* x, std::size_t n); | |
| void act_sigmoid_avx2(float* x, std::size_t n); | |
| void act_sigmoid_mul_avx2(float* o, const float* g, std::size_t n); | |
| float dot_int8_avx2(const std::int8_t* weights, const float* input, std::size_t n); | |
| float dot_int4_avx2(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const float* scales, std::size_t n); | |
| float dot_fp4_avx2(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const std::uint8_t* scales, std::size_t n); | |
| float dot_int8_q8_avx2(const std::int8_t* weights, const std::int16_t* act, const float* act_scales, std::size_t n); | |
| float dot_int4_q8_avx2(const std::uint8_t* weights, const std::int16_t* act, const float* scales, | |
| const float* act_scales, std::size_t n); | |
| float dot_fp4_q8_avx2(const std::uint8_t* weights, const std::int16_t* act, | |
| const std::uint8_t* scales, const float* act_scales, std::size_t n); | |
| // AVX-VNNI / AVX512-VNNI integer-MAC path (kernels_vnni.cpp; scalar-correct | |
| // fallback when the compiler lacks VNNI — dispatch never selects it then). | |
| float dot_int8_q8_vnni(const std::int8_t* weights, const std::int8_t* act, | |
| const float* act_scales, std::size_t n); | |
| // AVX512 zmm fp32 dot (kernels_vnni.cpp; needs AVX512F, gated on VNNI | |
| // presence like the rest of that TU). Different chunking/reduce from AVX2: | |
| // matches within fp32 rounding (PPL-gated), not bitwise-identical. | |
| float dot_fp32_avx512(const float* a, const float* b, std::size_t n); | |
| float dot_int8_avx512(const std::int8_t* weights, const float* input, std::size_t n); | |
| float dot_neon(const float* a, const float* b, std::size_t n); | |
| float dot_int8_neon(const std::int8_t* weights, const float* input, std::size_t n); | |
| float dot_int4_neon(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const float* scales, std::size_t n); | |
| float dot_fp4_neon(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const std::uint8_t* scales, std::size_t n); | |
| void permute_act32_neon(const float* input, std::size_t n, float* out); | |
| // SSE4.1 + AVX fallback TU (kernels_sse.cpp) for pre-AVX2 x86 | |
| // (Nehalem/Sandy/Bulldozer era). Same signatures as the AVX2 kernels; | |
| // dispatch prefers AVX2 whenever has_avx2_cpu() is true. | |
| float dot_sse41(const float* a, const float* b, std::size_t n); | |
| float dot_int8_sse41(const std::int8_t* weights, const float* input, std::size_t n); | |
| float dot_int4_sse41(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const float* scales, std::size_t n); | |
| float dot_fp4_sse41(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const std::uint8_t* scales, std::size_t n); | |
| float dot_avx_fp32(const float* a, const float* b, std::size_t n); | |
| } // namespace cism | |