Repository navigation
Expand file tree
/
Copy pathinkling.c
More file actions
4480 lines (4337 loc) · 218 KB
/
Copy pathinkling.c
File metadata and controls
4480 lines (4337 loc) · 218 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
/* Pure-C inference engine for Thinking Machines "Inkling" (text-only), Stage A.
* Goal, like olmoe.c before GLM-5.2: reproduce the EXACT token ids of the HF
* transformers reference (ref_inkling.json from tools/make_tiny_inkling.py)
* to validate the core math before scaling to the 975B checkpoint.
*
* Architecture (vs glm.c's MLA/RoPE/DSA — shares almost nothing):
* - hybrid attention: sliding-window layers (window=512, 16 KV heads) and
* global layers (8 KV heads) interleaved 5:1; conventional GQA, no RoPE
* - learned relative-position bias: r_proj(x) mixes a per-layer bank
* proj[d_rel, rel_extent] into one bias per backward distance
* - log-length scaling tau on global layers past n_floor tokens
* - depthwise-causal short convs (kernel 4, residual inside, fp32):
* on K and V inside attention, after attention, and after the MLP
* - MoE: sigmoid router + loss-free bias for top-k selection; combine
* weights are sigmoids of the raw logits jointly normalized over
* topk routed + n_shared shared experts, x route_scale x global_scale
* - logits: hidden / logits_mup_width_multiplier, sliced to unpadded vocab
*
* Dense weights (attn, norms, convs, router, shared experts, dense MLP)
* resident in RAM as f32; routed experts streamed from disk per-expert out
* of the fused [E, 2I, D] / [E, D, I] tensors, LRU-cached, optionally
* int-quantized (bits=0 keeps them f32 for bit-exact oracle validation).
*/
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <pthread.h>
#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__)
#include <sys/resource.h>
#include <sys/select.h> /* serve-loop stdin poll (POSIX); inkling serves on Linux */
#endif
#include "cli_args.h"
#include "st.h"
#include "tok.h"
#ifdef _OPENMP
#include <omp.h>
#endif
#include "omp_tune.h"
#include "route_trace.h"
#include "kv_prefix.h"
#include "pin_pool.h" /* KV prefix reuse (shared) */ /* shared routing telemetry (#700) */
#include "serve_codec.h"
#include "serve_budget.h"
#ifdef COLI_SEGMENT_ADAPTER
#include "segment_runtime.h"
#include "segment_adapters.h"
#include "segment_adapter_internal.h"
#endif
#ifdef COLI_EDGE_ADAPTER
#include "edge_runtime.h"
#include "edge_adapters.h"
#include "edge_tok_internal.h"
#endif
#ifdef COLI_CUDA
#include "backend_cuda_ink.h"
static int g_cuda = 0;
#endif
#ifdef COLI_METAL
/* Apple-GPU expert MoE (opt-in, COLI_METAL=1). Reuses colibri's batched
* coli_metal_moe_block: inkling's container int4 (nibble-packed, -8 offset,
* per-row scales) is bit-identical to the Metal fmt=2 kernel, and int8 to
* fmt=1. Expert slots live in page-aligned per-layer slabs registered once
* for zero-copy resolve — unified memory, no upload. Attention and the dense
* path stay on the CPU (at high hit rates ~90% of decode is expert matmul). */
#include "backend_metal.h"
static int g_metal = 0;
#endif
#ifdef COLI_VULKAN
/* Vulkan (opt-in, VK=1 build + COLI_VULKAN=1): the RESIDENT matrices, in the form
* they are kept in RAM, uploaded on first use -- the dense-int4g64 container's int8
* (fmt 1) and int4-g64 (fmt 4), and the f32 (fmt 10) and bf16 (fmt 11) residents of
* an unconverted snapshot -- and the routed experts through the shared expert tier
* (vk_tier.h, ink_vk_tier_start). The embedding row lookups stay on the CPU; CUDA
* keeps its own bf16 residents, and with CUDA (or Metal) on the tier stays off. */
#include "backend_vulkan.h"
#include "vk_tier.h"
#include "vk_chain.h"
static int g_vk_ready = 0;
static int g_vk_tier_try = 0; /* the device opened with the expert tier to be tried */
static int g_vk_chain = 0; /* COLI_VK_CHAIN decided on, and the chain's pipelines are up (inkling_chain.h) */
static int g_ink_dho = 0; /* COLI_VK_DENSE_HOST: the dense matrices on the device only (ink_dho_place) */
/* The partial chain's fit (inkc_fit_start, inkling_chain.h): the first n of L layers on the
* device; L = 0 while there is none (the chain off, or one that declines as before). */
static VkcFit g_inkc_fit;
#endif
#ifndef COLI_VULKAN /* exclusive RAM/VRAM (vk_tier.h) is the Vulkan build's: no device holds an expert */
static inline int vkt_ram_first(int layer, int eid) { (void)layer; (void)eid; return 0; }
static inline void vkt_ram_gave(void) {}
#endif
#if defined(__APPLE__)
#include <mach/mach.h>
#include <sys/sysctl.h>
#endif
#define MAXL 256
/* ---------- config ---------- */
typedef struct {
int hidden, n_layers, vocab, unpad_vocab;
int n_heads, n_kv, head_dim; /* global ("hybrid") layers */
int swa_heads, swa_kv, swa_hd; /* sliding ("hybrid_sliding") layers */
int window, d_rel, rel_extent, conv_k;
double log_floor; /* <=0: log scaling off */
float log_alpha;
int n_experts, topk, n_shared, moe_inter, dense_inter;
int eos;
float eps, route_scale, mup;
int audio_tok, mel_bins, mel_vocab; /* DMel audio input (TMLv0 <|audio|> placeholder) */
unsigned char local[MAXL]; /* 1 = sliding-window layer */
unsigned char sparse[MAXL]; /* 1 = MoE layer, 0 = dense MLP */
} Cfg;
/* per-layer dims that depend on the attention type */
#define L_HEADS(c,i) ((c)->local[i] ? (c)->swa_heads : (c)->n_heads)
#define L_KV(c,i) ((c)->local[i] ? (c)->swa_kv : (c)->n_kv)
#define L_HD(c,i) ((c)->local[i] ? (c)->swa_hd : (c)->head_dim)
#define L_EXT(c,i) ((c)->local[i] ? (c)->window : (c)->rel_extent)
/* ---------- resident weights ----------
* Large matmul weights keep their on-disk dtype in RAM: bf16 for the real
* 975B checkpoint (f32 residents would need ~172 GB, over sabre's 187),
* f32 for the tiny oracle (bit-exact validation). Under CUDA, bf16 tensors
* move to VRAM (dev set, host freed): decode reads ~35 GB of residents per
* token, so this trades the DDR5 bandwidth wall for VRAM bandwidth AND
* frees the same RAM for the expert cache. */
/* q4/qs: densa pre-quantizzata int4 group-scaled (gs=64) da un container
* separato, opzionale — vedi load_w. La densa bf16 di Inkling e' 49.4 GB e non
* entra in 25 GB (il caricamento la espande pure a f32, ~99 GB al picco: e' li'
* l'OOM). A int4-gs64 diventa ~15 GB. Nessun campo q4 = comportamento invariato. */
typedef struct { float *f; uint16_t *h; void *dev;
uint8_t *q4; /* int4 nibble-packed (qbits=4) o int8 (qbits=8) */
float *qs; /* scale: [rows*ng] se qbits=4, [rows] se qbits=8 */
int gs, qbits; int64_t qn; /* qn = byte in q4, per il guard OOB */
#ifdef COLI_VULKAN
void *vk; /* InkVk of the loaded tensor, shared by its wt_off_i
* views (NULL = CPU only) */
int64_t vk_vi
10BC8
ew; /* this view's first element in the tensor (wt_off_i):
* the key of its device copy, whatever the host holds */
const char *vk_name;/* the tensor's name, to read it back (ink_dho_reload) */
#endif
} Wt;
typedef struct {
float *in_ln, *post_ln;
Wt q, k, v, r, o; /* projections */
float *qn, *kn; /* per-head rmsnorm [head_dim] */
float *relp; /* [d_rel, ext] bias bank */
float *k_cw, *v_cw, *a_cw, *m_cw; /* sconv weights, [C*K] depthwise */
/* dense layers */
Wt dg, du, dd; float dgs;
/* MoE layers */
float *router, *rbias, rgs; /* [E+ns, D], [E], scalar */
Wt sh_g, sh_u, sh_d; /* shared experts [ns][I,D] etc. */
} Layer;
/* ---------- routed-expert cache: LRU + optional pinned set ----------
* Container snapshots keep the expert rows PACKED in RAM (int4 stays 4-bit:
* ~28 MB/expert instead of ~57 unpacked, so the same budget caches twice the
* experts); the matmul kernels unpack nibbles in-register. */
typedef struct {
int eid; uint64_t used;
int pinned; /* never evicted (usage-history pin) */
int filled; /* 0 while queued for a parallel fill */
uint8_t *p13, *p2; float *s13, *s2; /* container: packed rows + row scales */
int8_t *q13, *q2; /* bits>0: runtime-quantized int8 */
float *f13, *f2; /* bits==0: raw f32 (oracle) */
} Slot;
typedef struct {
Slot *slots;
int *slot_by_expert; /* expert id -> local slot, -1 if absent */
int n, cap;
} LCache;
/* One conversation's state for a multiplexed serve (KV_SLOTS>1, serve_mux below):
* its K and V (a sliding layer's ring), its four short-convolution banks, where it
* stands and the record of the tokens it holds. The Model holds the conversation a
* prefill runs on (ink_seq_swap trades it for a parked one); a multiplexed decode
* step parks them all and reads each row's from its InkRow. */
typedef struct { float **K, **V; float **cs[4]; int kv_len; kv_prefix kvp; } InkSeq;
typedef struct { InkSeq *seq; int pos; } InkRow;
static int g_ink_mux_slots = 1; /* KV_SLOTS: the conversations a serve decodes at once */
typedef struct {
Cfg c;
shards S;
shards Sq; /* container densa int4-gs64 (opzionale) */
int has_q, q_loaded; /* has_q: container presente; q_loaded: tensori presi da li */
int64_t q_bytes;
int quant_bits; /* 0 = f32 experts (oracle mode) */
int xq; /* experts on disk are a colibri container (U8 + .qs) */
Wt embed, lm_head;
float *embed_norm, *final_norm;
Wt audio_enc; /* [mel_bins*mel_vocab, D] embedding table */
float *audio_norm; /* audio tower RMSNorm [D]; NULL = no audio */
Layer *L;
LCache *cache;
int64_t rb13, rb2; /* container row-bytes (0 = not container) */
uint32_t **eusage; /* per-layer expert selection counts */
uint8_t **ehit; /* experts routed this turn, for HITS (dashboard Brain) */
int npin; /* pinned experts per sparse layer */
uint64_t clock, hits, miss;
uint64_t ereq, euse; /* routed richiesti (topk) vs usati dopo TOPP */
double t_fill, t_expert, t_shared, t_attn, t_route; /* phase timers */
float **K, **V; int kv_len, max_t; /* per-layer [kv][kv_ring_rows][hd]; sliding layers are a t%window ring */
float **cs[4]; /* conv states, [n_layers][C*(K-1)] */
double dense_load_s;
/* KV prefix reuse: what the current K/V and conv states were built from.
* See kv_prefix.h — recorded where the tokens are fed, never derived. */
kv_prefix kvp;
#ifdef COLI_VULKAN
void *vkchain; /* the dense chain's device state (inkling_chain.h), NULL until it runs */
void *vkchain2; /* its layers on COLI_VK_DEV2's device, after the primary's (inkling_chain.h) */
#endif
/* A multiplexed decode step (ink_step_rows): row s is the token at
* mux_rows[s].pos of mux_rows[s].seq; NULL in every other forward. */
const InkRow *mux_rows;
} Model;
/* ---------- utility ---------- */
static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; }
#if defined(__APPLE__)
static double rss_gb(void) { struct rusage r; getrusage(RUSAGE_SELF, &r); return r.ru_maxrss / (1024.0*1024.0*1024.0); }
#else
static double rss_gb(void) { struct rusage r; getrusage(RUSAGE_SELF, &r); return r.ru_maxrss / (1024.0*1024.0); }
#endif
static float *falloc(int64_t n) { float *p = malloc(n*sizeof(float)); if(!p){fprintf(stderr,"OOM %ld\n",(long)n);exit(1);} return p; }
static float sigmoidf(float x) { return 1.f / (1.f + expf(-x)); }
static float siluf(float x) { return x / (1.f + expf(-x)); }
/* TOPP=p (0..1): top-p adattivo sui routed — tieni gli esperti fino al peso
* cumulato p. 0 = spento (default): tutti i topk, calcolo invariato. */
static float g_topp = 0.f;
#include "matmul_f32.h" /* y[S,O] = x[S,I] @ W^T, W [O,I] f32 row-major */
#if defined(__AVX512BF16__) && defined(__AVX512F__)
#include <immintrin.h>
#define HAVE_BF16_DOT 1
#endif
#if defined(__AVX2__)
#include <immintrin.h>
#endif
/* bf16-weight matmul: activations rounded to bf16 per row (matches the HF
* bf16 reference numerics), hardware vdpbf16ps dot where available,
* shift-to-f32 scalar otherwise. */
static void matmul_h(float *y, const float *x, const uint16_t *W, int S, int I, int O) {
#ifdef HAVE_BF16_DOT
if (I % 32 == 0) {
uint16_t *xh = malloc((size_t)S * I * sizeof(uint16_t));
for (int s = 0; s < S; s++) {
const float *xs = x + (int64_t)s * I;
uint16_t *xd = xh + (int64_t)s * I;
for (int i = 0; i < I; i += 32) {
__m512 a = _mm512_loadu_ps(xs + i), b = _mm512_loadu_ps(xs + i + 16);
_mm512_storeu_si512(xd + i, (__m512i)_mm512_cvtne2ps_pbh(b, a));
}
}
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint16_t *w = W + (int64_t)o * I;
int s = 0;
for (; s + 3 < S; s += 4) {
const uint16_t *x0 = xh + (int64_t)(s+0)*I;
const uint16_t *x1 = xh + (int64_t)(s+1)*I;
const uint16_t *x2 = xh + (int64_t)(s+2)*I;
const uint16_t *x3 = xh + (int64_t)(s+3)*I;
__m512 a0 = _mm512_setzero_ps(), a1 = a0, a2 = a0, a3 = a0;
for (int i = 0; i < I; i += 32) {
__m512bh wb = (__m512bh)_mm512_loadu_si512(w + i);
a0 = _mm512_dpbf16_ps(a0, (__m512bh)_mm512_loadu_si512(x0 + i), wb);
a1 = _mm512_dpbf16_ps(a1, (__m512bh)_mm512_loadu_si512(x1 + i), wb);
a2 = _mm512_dpbf16_ps(a2, (__m512bh)_mm512_loadu_si512(x2 + i), wb);
a3 = _mm512_dpbf16_ps(a3, (__m512bh)_mm512_loadu_si512(x3 + i), wb);
}
y[(int64_t)(s+0)*O + o] = _mm512_reduce_add_ps(a0);
y[(int64_t)(s+1)*O + o] = _mm512_reduce_add_ps(a1);
y[(int64_t)(s+2)*O + o] = _mm512_reduce_add_ps(a2);
y[(int64_t)(s+3)*O + o] = _mm512_reduce_add_ps(a3);
}
for (; s < S; s++) {
const uint16_t *xs = xh + (int64_t)s*I;
__m512 acc = _mm512_setzero_ps();
for (int i = 0; i < I; i += 32)
acc = _mm512_dpbf16_ps(acc, (__m512bh)_mm512_loadu_si512(xs + i),
(__m512bh)_mm512_loadu_si512(w + i));
y[(int64_t)s*O + o] = _mm512_reduce_add_ps(acc);
}
}
free(xh);
return;
}
#endif
/* Prefill tile: decode each bf16 weight once for four independent rows.
* Explicit mul+add (rather than a C expression that the compiler may fuse)
* matches the non-FMA scalar oracle bit for bit. Every SIMD lane still
* accumulates i=0..I-1 in the original order. */
#if defined(__AVX2__)
if (S > 1) {
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint16_t *w = W + (int64_t)o * I;
int s = 0;
for (; s + 3 < S; s += 4) {
const float *x0 = x + (int64_t)(s+0)*I;
const float *x1 = x + (int64_t)(s+1)*I;
const float *x2 = x + (int64_t)(s+2)*I;
const float *x3 = x + (int64_t)(s+3)*I;
__m128 acc = _mm_setzero_ps();
for (int i = 0; i < I; i++) {
union { uint32_t u; float f; } v = { (uint32_t)w[i] << 16 };
__m128 xv = _mm_set_ps(x3[i], x2[i], x1[i], x0[i]);
acc = _mm_add_ps(acc, _mm_mul_ps(xv, _mm_set1_ps(v.f)));
}
float a[4]; _mm_storeu_ps(a, acc);
y[(int64_t)(s+0)*O + o] = a[0]; y[(int64_t)(s+1)*O + o] = a[1];
y[(int64_t)(s+2)*O + o] = a[2]; y[(int64_t)(s+3)*O + o] = a[3];
}
for (; s < S; s++) {
const float *xs = x + (int64_t)s * I;
float acc = 0.f;
for (int i = 0; i < I; i++) {
union { uint32_t u; float f; } v = { (uint32_t)w[i] << 16 };
acc += xs[i] * v.f;
}
y[(int64_t)s * O + o] = acc;
}
}
return;
}
#endif
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint16_t *w = W + (int64_t)o * I;
for (int s = 0; s < S; s++) {
const float *xs = x + (int64_t)s * I;
float acc = 0.f;
for (int i = 0; i < I; i++) {
union { uint32_t u; float f; } v = { (uint32_t)w[i] << 16 };
acc += xs[i] * v.f;
}
y[(int64_t)s * O + o] = acc;
}
}
}
/* dispatch on where the weight lives */
/* y[S,O] = x[S,I] @ W^T con W int4 GROUP-scaled: nibble +8, low = colonna pari,
* una scala f32 ogni `gs` elementi lungo I (ng = ceil(I/gs) scale per riga).
* Differenza da matmul_q4 (per-riga): la scala cambia DENTRO la riga, quindi
* l'accumulo va chiuso a ogni gruppo invece che una volta sola a fine riga. */
#if defined(__AVX2__) && defined(__FMA__)
static inline float matmul_i4g_row_fma(const float *xs, const uint8_t *w,
const float *sc, int I, int gs, int ng) {
__m128 acc = _mm_setzero_ps();
for (int g = 0; g < ng; g++) {
int i0 = g*gs, i1 = i0+gs; if (i1 > I) i1 = I;
__m128 part = _mm_setzero_ps();
for (int i = i0; i < i1; i++) {
uint8_t b = w[i >> 1]; int q = (i & 1) ? (b >> 4) : (b & 15);
part = _mm_fmadd_ss(_mm_set_ss(xs[i]), _mm_set_ss((float)(q-8)), part);
}
acc = _mm_add_ss(acc, _mm_mul_ss(part, _mm_set_ss(sc[g])));
}
return _mm_cvtss_f32(acc);
}
#endif
static void matmul_i4g(float *y, const float *x, const uint8_t *p, const float *scale,
int S, int I, int O, int gs) {
int ng = (I + gs - 1) / gs;
int64_t rb = (I + 1) / 2;
#if defined(__AVX2__) && defined(__FMA__)
/* Four prompt rows share nibble unpack and scale loads. The explicit FMA
* for each lane and the group-boundary mul+add match the scalar kernel's
* generated operation order bit for bit. */
if (S > 1) {
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint8_t *w = p + (int64_t)o * rb;
const float *sc = scale + (int64_t)o * ng;
int s = 0;
for (; s + 3 < S; s += 4) {
const float *x0 = x + (int64_t)(s+0)*I, *x1 = x + (int64_t)(s+1)*I;
const float *x2 = x + (int64_t)(s+2)*I, *x3 = x + (int64_t)(s+3)*I;
__m128 acc = _mm_setzero_ps();
for (int g = 0; g < ng; g++) {
int i0 = g*gs, i1 = i0+gs; if (i1 > I) i1 = I;
__m128 part = _mm_setzero_ps();
for (int i = i0; i < i1; i++) {
uint8_t b = w[i >> 1]; int q = (i & 1) ? (b >> 4) : (b & 15);
__m128 xv = _mm_set_ps(x3[i], x2[i], x1[i], x0[i]);
part = _mm_fmadd_ps(xv, _mm_set1_ps((float)(q-8)), part);
}
acc = _mm_add_ps(acc, _mm_mul_ps(part, _mm_set1_ps(sc[g])));
}
float a[4]; _mm_storeu_ps(a, acc);
y[(int64_t)(s+0)*O + o] = a[0]; y[(int64_t)(s+1)*O + o] = a[1];
y[(int64_t)(s+2)*O + o] = a[2]; y[(int64_t)(s+3)*O + o] = a[3];
}
for (; s < S; s++) {
const float *xs = x + (int64_t)s*I;
y[(int64_t)s*O + o] = matmul_i4g_row_fma(xs, w, sc, I, gs, ng);
}
}
return;
}
#endif
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint8_t *w = p + (int64_t)o * rb;
const float *sc = scale + (int64_t)o * ng;
for (int s = 0; s < S; s++) {
const float *xs = x + (int64_t)s * I;
#if defined(__AVX2__) && defined(__FMA__)
y[(int64_t)s * O + o] = matmul_i4g_row_fma(xs, w, sc, I, gs, ng);
#else
float acc = 0.f;
for (int g = 0; g < ng; g++) {
int i0 = g * gs, i1 = i0 + gs; if (i1 > I) i1 = I;
float part = 0.f;
for (int i = i0; i < i1; i++) {
uint8_t b = w[i >> 1];
int q = (i & 1) ? (b >> 4) : (b & 0x0F);
part += xs[i] * (float)(q - 8);
}
acc += part * sc[g]; /* scala chiusa per gruppo */
}
y[(int64_t)s * O + o] = acc;
#endif
}
}
}
/* int8 per-riga (embed/lm_head: sensibili, non vanno a 4 bit) */
static void matmul_i8r(float *y, const float *x, const int8_t *q, const float *scale,
int S, int I, int O) {
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const int8_t *w = q + (int64_t)o * I; float sc = scale[o];
for (int s = 0; s < S; s++) {
const float *xs = x + (int64_t)s * I;
float acc = 0.f;
for (int i = 0; i < I; i++) acc += xs[i] * (float)w[i];
y[(int64_t)s * O + o] = acc * sc;
}
}
}
#ifdef COLI_VULKAN
/* Device copies of one loaded tensor. A Wt travels by value and the shared
* experts are read through wt_off_i views of one fused [ns][R,I] tensor, so the
* cache cannot live in the copy matmul_w receives: it lives in a table the
* loaded Wt points to and every view copies, one entry per view (keyed by the
* view's first weight). Released with the tensor it belongs to. */
typedef struct { int64_t view; int used; ColiVkTensor *t; int dead; } InkVkView;
/* gone: the device holds the tensor alone (COLI_VK_DENSE_HOST): fmt and gs are the ones
* its views went up in, home the loaded Wt the CPU reads it back into. */
typedef struct { int n, fmt, gs, gone; Wt *home; InkVkView e[]; } InkVk;
/* The shader format that reads a weight exactly as the CPU kernel does, 0 for none.
* int8 per row (matmul_i8r) is fmt 1; int4 group-scaled with +8 nibbles and the low
* nibble on the even column (matmul_i4g) is fmt 4; f32 (matmul) is fmt 10; bf16
* (matmul_h) is fmt 11 -- except where matmul_h rounds the activations to bf16
* (the AVX512-BF16 dot), which the shader does not: there bf16 stays on the CPU. */
static int ink_vk_fmt(const Wt *w) {
if (w->q4) return w->qbits == 8 ? 1 : (w->qbits == 4 && w->gs >= 8 && w->gs % 8 == 0) ? 4 : 0;
if (w->f) return 10;
#ifndef HAVE_BF16_DOT
if (w->h) return 11;
#endif
return 0;
}
/* Attach a table to a loaded tensor the shader can read; returns its fmt, or 0. */
static int ink_vk_attach(Wt *w, int views) {
int fmt = ink_vk_fmt(w);
if (!fmt || w->vk || views < 1) return 0;
InkVk *v = calloc(1, sizeof(*v) + (size_t)views * sizeof(v->e[0]));
if (!v) return 0;
v->n = views; v->fmt = fmt; v->gs = w->q4 ? w->gs : 0; v->home = w;
w->vk = v;
return fmt;
}
static void ink_vk_release(Wt *w) {
InkVk *v = w->vk;
if (!v) return;
for (int k = 0; k < v->n; k++) if (v->e[k].t) coli_vk_tensor_free(v->e[k].t);
free(v);
w->vk = NULL;
}
/* 1 when y was computed on the device. One command buffer in the backend:
* main thread only, never from inside a parallel region. */
/* The entry of a view (keyed by its first element in the tensor), claimed on first use;
* NULL when every entry belongs to another view. */
static InkVkView *ink_vk_entry(InkVk *v, int64_t view) {
for (int k = 0; k < v->n; k++) {
InkVkView *e = &v->e[k];
if (e->used && e->view != view) continue;
e->used = 1; e->view = view;
return e;
}
return NULL;
}
static int ink_vk_matmul(float *y, const float *x, const Wt *W, int S, int I, int O) {
InkVk *v = W->vk;
if (!g_vk_ready || !v) return 0;
#ifdef _OPENMP
if (omp_in_parallel()) return 0;
#endif
int fmt = v->gone ? v->fmt : ink_vk_fmt(W);
const void *data = W->q4 ? (const void *)W->q4 : W->f ? (const void *)W->f : (const void *)W->h;
if (!fmt || (!data && !v->gone)) return 0;
InkVkView *e = ink_vk_entry(v, W->vk_view);
if (!e || e->dead || (v->gone && !e->t)) return 0; /* gone: only a copy made before the drop serves */
if (coli_vk_matmul(&e->t, y, x, data, W->q4 ? W->qs : NULL, fmt, S, I, O, v->gs))
return 1;
if (!e->t) e->dead = 1; /* refused at upload: this view stays on the CPU */
return 0;
}
/* A view's device copy, made now if it is not there yet (the chain's tensors and the
* eager upload of COLI_VK_DENSE_HOST); NULL when the view has none. */
static ColiVkTensor *ink_vk_view_tensor(const Wt *W, int I, int O) {
InkVk *v = W->vk;
if (!v) return NULL;
int fmt = v->gone ? v->fmt : ink_vk_fmt(W);
const void *data = W->q4 ? (const void *)W->q4 : W->f ? (const void *)W->f : (const void *)W->h;
if (!fmt || (!data && !v->gone)) return NULL;
InkVkView *e = ink_vk_entry(v, W->vk_view);
if (!e || e->dead) return NULL;
if (!e->t && (v->gone || !coli_vk_tensor_ensure(&e->t, data, W->q4 ? W->qs : NULL, fmt, I, O, v->gs))) {
e->dead = 1; return NULL;
}
return e->t;
}
static Wt ink_dho_reload(Wt W, int I); /* below, beside load_w */
static void ink_vk_report(void) {
if (g_vk_ready)
fprintf(stderr, "[VK] inkling: %llu matmuls on the GPU\n", coli_vk_matmul_calls());
}
#endif
#ifdef COLI_INKLING_SHARED_BATCH_TEST
static uint64_t g_matmul_w_calls;
#endif
static void matmul_w(float *y, const float *x, Wt W, int S, int I, int O) {
#ifdef COLI_INKLING_SHARED_BATCH_TEST
g_matmul_w_calls++;
#endif
#ifdef COLI_VULKAN
if (W.vk && ((InkVk *)W.vk)->gone) { /* the device holds it alone */
if (ink_vk_matmul(y, x, &W, S, I, O)) return;
W = ink_dho_reload(W, I); /* the CPU needs it after all (a lost device): from disk */
}
#endif
#ifdef COLI_CUDA
if (W.dev) {
if (ink_cuda_matmul_bf16(y, x, W.dev, S, I, O) == 0) return;
fprintf(stderr, "cuda matmul failed and host copy was freed\n"); exit(1);
}
#endif
if (W.q4) {
/* guard: il container e' un file, non un invariante — se la geometria non
* torna si esce invece di leggere fuori dal buffer. */
int64_t need = W.qbits == 8 ? (int64_t)O * I : (int64_t)O * ((I + 1) / 2);
if (need > W.qn) {
fprintf(stderr, "dense q4: geometria incoerente (serve %lld B, ho %lld) I=%d O=%d\n",
(long long)need, (long long)W.qn, I, O); exit(1);
}
#ifdef COLI_VULKAN
if (W.vk && ink_vk_matmul(y, x, &W, S, I, O)) return;
#endif
if (W.qbits == 8) matmul_i8r(y, x, (const int8_t*)W.q4, W.qs, S, I, O);
else matmul_i4g(y, x, W.q4, W.qs, S, I, O, W.gs);
return;
}
#ifdef COLI_VULKAN
if (W.vk && ink_vk_matmul(y, x, &W, S, I, O)) return;
#endif
if (W.f) matmul(y, x, W.f, S, I, O);
else matmul_h(y, x, W.h, S, I, O);
}
/* y[1,O] = x @ q^T, int8 weights + per-row scale. Fast path: activations
* quantized Q8 per 32-block, VNNI (or maddubs) int8 dot — same family as
* glm.c's IDOT kernels; IDOT=0 falls back to the byte-exact scalar route. */
#if defined(__AVX2__)
static inline __m256i i8dot_block(__m256i acc, __m256i a, __m256i b) {
__m256i ax = _mm256_sign_epi8(a, a); /* |a| as u8 */
__m256i sy = _mm256_sign_epi8(b, a); /* b * sign(a) */
#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
return _mm256_dpbusd_epi32(acc, ax, sy);
#else
__m256i p = _mm256_maddubs_epi16(ax, sy);
return _mm256_add_epi32(acc, _mm256_madd_epi16(p, _mm256_set1_epi16(1)));
#endif
}
#endif
static void matmul_q(float *y, const float *x, const int8_t *q, const float *scale, int I, int O) {
#if defined(__AVX2__)
static int idot = -1;
/* Opt-in (IDOT=1), not default: the fast path only exists under __AVX2__
* and quantizes ACTIVATIONS per 32-block — the same model produced
* different tokens on x86 vs ARM with the old on-by-default. Same class
* and same fix as olmoe (#1044) and qwen36 (#712 review). */
if (idot < 0) { const char *e = getenv("IDOT"); idot = (e && atoi(e)); }
if (idot && I % 32 == 0 && I <= 8192) {
int nb = I / 32;
int8_t xi[8192]; float xs[256];
for (int b = 0; b < nb; b++) {
const float *xb = x + b*32;
float am = 0.f; for (int i = 0; i < 32; i++) { float a = fabsf(xb[i]); if (a > am) am = a; }
float s = am/127.f; if (s < 1e-12f) s = 1e-12f;
xs[b] = s; float inv = 1.f/s;
for (int i = 0; i < 32; i++) xi[b*32+i] = (int8_t)lrintf(xb[i]*inv);
}
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const int8_t *w = q + (int64_t)o * I;
float acc = 0.f;
for (int b = 0; b < nb; b++) {
__m256i vacc = i8dot_block(_mm256_setzero_si256(),
_mm256_loadu_si256((const __m256i*)(xi + b*32)),
_mm256_loadu_si256((const __m256i*)(w + b*32)));
__m128i lo = _mm256_castsi256_si128(vacc), hi = _mm256_extracti128_si256(vacc, 1);
__m128i s4 = _mm_add_epi32(lo, hi);
s4 = _mm_hadd_epi32(s4, s4); s4 = _mm_hadd_epi32(s4, s4);
acc += xs[b] * (float)_mm_cvtsi128_si32(s4);
}
y[o] = acc * scale[o];
}
return;
}
#endif
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const int8_t *w = q + (int64_t)o * I;
float acc = 0.f;
for (int i = 0; i < I; i++) acc += x[i] * (float)w[i];
y[o] = acc * scale[o];
}
}
/* y[1,O] = x @ W^T with W kept PACKED int4 (low nibble = even column, +8
* offset, per-row scale — the on-disk container layout, cached as-is).
* Nibbles unpack in-register: same numeric result as unpack-to-int8 +
* matmul_q, half the cache footprint. IDOT=0 keeps the byte-exact scalar. */
static void matmul_q4(float *y, const floa
10BC8
t *x, const uint8_t *p, const float *scale, int I, int O) {
#if defined(__AVX2__)
static int idot = -1;
/* Opt-in (IDOT=1), not default: the fast path only exists under __AVX2__
* and quantizes ACTIVATIONS per 32-block — the same model produced
* different tokens on x86 vs ARM with the old on-by-default. Same class
* and same fix as olmoe (#1044) and qwen36 (#712 review). */
if (idot < 0) { const char *e = getenv("IDOT"); idot = (e && atoi(e)); }
if (idot && I % 32 == 0 && I <= 8192) {
int nb = I / 32;
int8_t xi[8192]; float xs[256];
for (int b = 0; b < nb; b++) {
const float *xb = x + b*32;
float am = 0.f; for (int i = 0; i < 32; i++) { float a = fabsf(xb[i]); if (a > am) am = a; }
float s = am/127.f; if (s < 1e-12f) s = 1e-12f;
xs[b] = s; float inv = 1.f/s;
for (int i = 0; i < 32; i++) xi[b*32+i] = (int8_t)lrintf(xb[i]*inv);
}
const __m128i m4 = _mm_set1_epi8(0x0F);
const __m256i b8 = _mm256_set1_epi8(8);
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint8_t *w = p + (int64_t)o * (I/2);
float acc = 0.f;
for (int b = 0; b < nb; b++) {
__m128i by = _mm_loadu_si128((const __m128i*)(w + b*16)); /* 16 B = 32 nibbles */
__m128i lo = _mm_and_si128(by, m4); /* even columns */
__m128i hi = _mm_and_si128(_mm_srli_epi16(by, 4), m4); /* odd columns */
__m256i nib = _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), /* cols 16..31 */
_mm_unpacklo_epi8(lo, hi)); /* cols 0..15 */
nib = _mm256_sub_epi8(nib, b8);
__m256i vacc = i8dot_block(_mm256_setzero_si256(),
_mm256_loadu_si256((const __m256i*)(xi + b*32)), nib);
__m128i l = _mm256_castsi256_si128(vacc), h = _mm256_extracti128_si256(vacc, 1);
__m128i s4 = _mm_add_epi32(l, h);
s4 = _mm_hadd_epi32(s4, s4); s4 = _mm_hadd_epi32(s4, s4);
acc += xs[b] * (float)_mm_cvtsi128_si32(s4);
}
y[o] = acc * scale[o];
}
return;
}
#endif
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const uint8_t *w = p + (int64_t)o * (I/2);
float acc = 0.f;
for (int i = 0; i < I; i += 2) {
uint8_t byte = w[i/2];
acc += x[i] * (float)((int)(byte & 0xF) - 8);
acc += x[i+1] * (float)((int)(byte >> 4) - 8);
}
y[o] = acc * scale[o];
}
}
static void quantize_rows(const float *w, int8_t *q, float *scale, int O, int I, int bits) {
int qmax = (1 << (bits - 1)) - 1;
#pragma omp parallel for schedule(static)
for (int o = 0; o < O; o++) {
const float *wr = w + (int64_t)o * I;
float amax = 0.f; for (int i = 0; i < I; i++) { float a = fabsf(wr[i]); if (a > amax) amax = a; }
float s = amax / qmax; if (s < 1e-8f) s = 1e-8f;
scale[o] = s;
int8_t *qr = q + (int64_t)o * I;
for (int i = 0; i < I; i++) {
int v = (int)lrintf(wr[i] / s);
if (v > qmax) v = qmax;
if (v < -qmax-1) v = -qmax-1;
qr[i] = (int8_t)v;
}
}
}
/* rmsnorm computed in f64 accumulate like the f32->f32 reference */
static void rmsnorm_row(float *out, const float *x, const float *w, int D, float eps) {
double ms = 0; for (int i = 0; i < D; i++) ms += (double)x[i]*x[i];
float r = 1.f / sqrtf((float)(ms / D) + eps);
for (int i = 0; i < D; i++) out[i] = x[i] * r * w[i];
}
static void softmax_row(float *x, int n) {
float m = -1e30f; for (int i = 0; i < n; i++) if (x[i] > m) m = x[i];
float s = 0; for (int i = 0; i < n; i++) { x[i] = expf(x[i]-m); s += x[i]; }
for (int i = 0; i < n; i++) x[i] /= s;
}
/* ---------- depthwise causal short conv, residual inside (fp32) ----------
* seq[S,C] in-place: out[t] = sum_j w[c,j]*in[t+j-(K-1)] + in[t], history from
* state[C*(K-1)] (raw pre-conv inputs), which is updated to the new tail. */
static void sconv_apply(float *seq, int S, int C, const float *w, float *state, int K) {
int P = K - 1;
#pragma omp parallel
{
float *col = malloc((P + S) * sizeof(float));
#pragma omp for schedule(static)
for (int ch = 0; ch < C; ch++) {
for (int j = 0; j < P; j++) col[j] = state[(int64_t)ch*P + j];
for (int t = 0; t < S; t++) col[P + t] = seq[(int64_t)t*C + ch];
const float *wc = w + (int64_t)ch*K;
for (int t = 0; t < S; t++) {
float acc = 0.f;
for (int j = 0; j < K; j++) acc += wc[j] * col[t + j];
seq[(int64_t)t*C + ch] = acc + col[P + t];
}
for (int j = 0; j < P; j++) state[(int64_t)ch*P + j] = col[S + j];
}
free(col);
}
}
/* ---------- config loading ----------
* Accepts both the flat text config (tiny oracle via InklingForCausalLM) and
* the full multimodal config.json (real checkpoint, fields under text_config). */
static double jnum(jval *o, const char *k, double dflt) {
jval *v = json_get(o, k);
return (v && v->t == J_NUM) ? v->num : dflt;
}
static void load_cfg(Cfg *c, const char *snap) {
char path[2048]; snprintf(path, sizeof(path), "%s/config.json", snap);
FILE *f = fopen(path, "rb"); if(!f){perror(path);exit(1);}
fseek(f,0,SEEK_END); long n=ftell(f); fseek(f,0,SEEK_SET);
char *buf = malloc(n+1); if(fread(buf,1,n,f)!=(size_t)n){} buf[n]=0; fclose(f);
char *arena=NULL; jval *root = json_parse(buf, &arena);
jval *r = json_get(root, "text_config"); if (!r) r = root;
c->hidden = (int)jnum(r,"hidden_size",6144);
c->n_layers = (int)jnum(r,"num_hidden_layers",66);
c->vocab = (int)jnum(r,"vocab_size",201024);
c->unpad_vocab = (int)jnum(r,"unpadded_vocab_size",c->vocab);
c->n_heads = (int)jnum(r,"num_attention_heads",64);
c->n_kv = (int)jnum(r,"num_key_value_heads",8);
c->head_dim = (int)jnum(r,"head_dim",128);
c->swa_heads = (int)jnum(r,"swa_num_attention_heads",c->n_heads);
c->swa_kv = (int)jnum(r,"swa_num_key_value_heads",16);
c->swa_hd = (int)jnum(r,"swa_head_dim",c->head_dim);
c->window = (int)jnum(r,"sliding_window_size",512);
c->d_rel = (int)jnum(r,"d_rel",16);
c->rel_extent = (int)jnum(r,"rel_extent",1024);
c->log_floor = jnum(r,"log_scaling_n_floor",0);
c->log_alpha = (float)jnum(r,"log_scaling_alpha",0.1);
c->conv_k = (int)jnum(r,"sconv_kernel_size", jnum(r,"conv_kernel_size",4));
c->n_experts = (int)jnum(r,"n_routed_experts",256);
c->topk = (int)jnum(r,"num_experts_per_tok",6);
c->n_shared = (int)jnum(r,"n_shared_experts",2);
c->eps = (float)jnum(r,"rms_norm_eps",1e-6);
c->route_scale = (float)jnum(r,"route_scale",8.0);
c->mup = (float)jnum(r,"logits_mup_width_multiplier",24.0);
/* eos lives at the top level in the real multimodal config, in the text
* config for a flat snapshot; may be null (tiny oracle) */
jval *eo = json_get(root,"eos_token_id");
if (!eo || eo->t != J_NUM) eo = json_get(r,"eos_token_id");
c->eos = (eo && eo->t == J_NUM) ? (int)eo->num : -1;
/* DMel audio: placeholder id at the top level (absent in the shipped
* config -> the TMLv0 constant), frame geometry under audio_config */
c->audio_tok = (int)jnum(root,"audio_token_id",200023);
jval *ac = json_get(root,"audio_config");
c->mel_bins = ac ? (int)jnum(ac,"n_mel_bins",80) : 80;
c->mel_vocab = ac ? (int)jnum(ac,"mel_vocab_size",16) : 16;
/* real config.json: intermediate_size = MoE, dense_intermediate_size = dense.
* HF-saved config (post_init applied): intermediate_size = dense, moe_intermediate_size = MoE. */
jval *dis = json_get(r,"dense_intermediate_size");
if (dis && dis->t == J_NUM) {
c->dense_inter = (int)dis->num;
c->moe_inter = (int)jnum(r,"intermediate_size",3072);
} else {
c->dense_inter = (int)jnum(r,"intermediate_size",24576);
c->moe_inter = (int)jnum(r,"moe_intermediate_size",3072);
}
if (c->n_layers > MAXL) { fprintf(stderr,"n_layers %d > MAXL\n", c->n_layers); exit(1); }
/* SEC: these loops run to n_layers -- a number from config.json, up to
* MAXL -- while indexing a JSON array whose length is a different, equally
* attacker-chosen number. "num_hidden_layers": 66 with a one-element
* layer_types[] read kids[1..65] past a malloc'd array that starts at
* capacity 8, then dereferenced whatever was there as a string. config.json
* alone was enough; no weights and no valid snapshot were needed.
*
* A short array now means "not specified for these layers", so the loop
* falls through to the same default it would have used had the key been
* absent. The element type is checked too: kids[i]->str is a union member
* that is only a valid pointer when t == J_STR. */
#define LT_STR(arr, i) \
((arr) && (arr)->t == J_ARR && (i) < (arr)->len && \
(arr)->kids[i] && (arr)->kids[i]->t == J_STR ? (arr)->kids[i]->str : NULL)
/* attention layer types: explicit layer_types[] > local_layer_ids[] > (i+1)%6 rule */
jval *lt = json_get(r,"layer_types");
jval *ll = json_get(r,"local_layer_ids");
for (int i = 0; i < c->n_layers; i++) {
const char *ltype = LT_STR(lt, i);
if (ltype) c->local[i] = (strcmp(ltype,"hybrid_sliding")==0);
else if (ll && ll->t == J_ARR) {
c->local[i] = 0;
for (int j = 0; j < ll->len; j++)
if (ll->kids[j] && ll->kids[j]->t == J_NUM &&
(int)ll->kids[j]->num == i) { c->local[i] = 1; break; }
} else c->local[i] = ((i + 1) % 6) != 0;
}
/* MLP types: explicit mlp_layer_types[] > dense_mlp_idx (first k layers dense) */
jval *mt = json_get(r,"mlp_layer_types");
int dense_idx = (int)jnum(r,"dense_mlp_idx",0);
for (int i = 0; i < c->n_layers; i++) {
const char *mtype = LT_STR(mt, i);
if (mtype) c->sparse[i] = (strcmp(mtype,"sparse")==0);
else c->sparse[i] = (i >= dense_idx);
}
#undef LT_STR
free(buf); free(arena);
}
/* ---------- weight loading ---------- */
/* The embedding norm has two names. The converter and every container written with
* transformers up to 5.17 call it model.embed_norm.weight; transformers 5.18 saves
* the same tensor as model.embed_tokens.embed_norm.weight. Read by the old name
* only, a checkpoint saved by 5.18 loaded with no embedding norm at all -- no
* error, and every answer wrong (the tiny oracle: 1 of 36 positions). */
static const char *embed_norm_name(shards *S) {
if (st_has(S, "model.embed_norm.weight")) return "model.embed_norm.weight";
if (st_has(S, "model.embed_tokens.embed_norm.weight")) return "model.embed_tokens.embed_norm.weight";
return NULL;
}
static float *load_t(Model *m, const char *name) {
int64_t n = st_numel(&m->S, name);
if (n < 0) { fprintf(stderr, "missing %s\n", name); exit(1); }
float *p = falloc(n);
st_read_f32(&m->S, name, p, 0);
return p;
}
static float load_scalar(Model *m, const char *name, float dflt) {
if (!st_has(&m->S, name)) return dflt;
/* SEC: `v` is four bytes on the stack, and st_read_f32 writes as many as
* the FILE says the tensor holds -- it validates the header against itself
* (numel*esz == nbytes) and cannot know the destination size. A crafted
* snapshot declaring e.g. model.layers.0.mlp.gate.global_scale as F32
* shape [4096] therefore drops 16 KiB of attacker bytes over this frame
* and its return address; the reported crash had RSP fully controlled.
*
* st_read_f32_cap is the same read with the caller's capacity passed in.
* A scalar's capacity is 1. Anything larger is a hostile or corrupt file
* and stops the load, which is the behaviour every other bounds check in
* st.h already has. */
float v; st_read_f32_cap(&m->S, name, &v, 1, 0); return v;
}
/* chunked pread: a single pread caps at ~2.1 GB on Linux, and the bf16
* embed/lm_head tensors are 2.47 GB — loop in 1 GB slices *
9D75
/
static void pread_all(int fd, void *buf, int64_t nb, int64_t off) {
char *p = buf;
while (nb > 0) {
int64_t chunk = nb < (1<<30) ? nb : (1<<30);
ssize_t got = pread(fd, p, (size_t)chunk, off);
if (got <= 0) { perror("pread chunk"); exit(1); }
p += got; off += got; nb -= got;
}
}
/* big matmul weights keep their on-disk dtype resident: BF16 raw (real
* checkpoint, halves RAM), anything else as f32 (tiny oracle: bit-exact).
* gpu_ok: bf16 tensors move to VRAM while budget lasts (embed stays host —
* it's a row lookup, not a matmul). */
/* Container densa pre-quantizzata (opzionale): <snap>/dense-int4g64/.
* Se il tensore c'e' li' dentro lo carichiamo int4-gs64 (U8 + sidecar .qs, la
* stessa convenzione che il container usa gia' per gli esperti); altrimenti si
* prosegue col percorso bf16/f32 di sempre. Container assente = nessun cambio. */
static int load_w_quant(Model *m, const char *name, int64_t orig_numel, Wt *out) {
if (!m->has_q) return 0;
st_tensor *t = st_find(&m->Sq, name);
if (!t || t->dtype != 3) return 0; /* 3 = U8/byte grezzi */
char qn[352]; snprintf(qn, sizeof(qn), "%s.qs", name);
st_tensor *s = st_find(&m->Sq, qn);
if (!s) return 0; /* senza scale non si decodifica */
/* Geometria dedotta dai CONTEGGI, non dalle shape (st_tensor non le porta):
* int4-gs64 -> byte ~= numel/2 e scale ~= numel/64
* int8 -> byte == numel e scale == righe (numel/I, I ignoto qui:
* basta che le scale siano molte meno dei byte)
* Il controllo forte sull'OOB lo fa matmul_w con l'I vero del chiamante. */
Wt w = {0};
if (t->nbytes == orig_numel && s->numel * 64 < (int64_t)t->nbytes) {
w.qbits = 8; /* int8 per riga */
} else if (t->nbytes * 2 >= orig_numel && t->nbytes * 2 <= orig_numel + 2 * s->numel) {
w.qbits = 4; w.gs = 64; /* int4 group-scaled */
} else {
fprintf(stderr, "[dense] %s: geometria non riconosciuta nel container, uso il bf16\n", name);
return 0; /* dubbio -> percorso originale */
}
w.q4 = malloc(t->nbytes); if (!w.q4) { fprintf(stderr, "OOM %s\n", name); exit(1); }
pread_all(t->fd, w.q4, t->nbytes, t->off);
w.qn = t->nbytes;
w.qs = falloc(s->numel);
st_read_f32(&m->Sq, qn, w.qs, 0);
*out = w;
m->q_loaded++; m->q_bytes += t->nbytes + (int64_t)s->numel * 4;
return 1;
}
#ifdef COLI_VULKAN
/* the names of loaded tensors, kept for the life of the process (a Wt travels by value) */
static const char *ink_name_keep(const char *name) {
char *n = strdup(name);
if (!n) { fprintf(stderr, "OOM weight name\n"); exit(1); }
return n;
}
#endif
static Wt load_w(Model *m, const char *name, int gpu_ok) {
Wt w = {0};
st_tensor *t = st_find(&m->S, name);
if (!t) { fprintf(stderr, "missing %s\n", name); exit(1); }
#ifdef COLI_VULKAN
if (load_w_quant(m, name, t->numel, &w)) { w.vk_name = ink_name_keep(name); return w; }
w.vk_name = ink_name_keep(name);
#else
if (load_w_quant(m, name, t->numel, &w)) return w;
#endif
if (t->dtype == 0) {
w.h = malloc(t->nbytes); if (!w.h) { fprintf(stderr,"OOM %s\n",name); exit(1); }
pread_all(t->fd, w.h, t->nbytes, t->off);
#ifdef COLI_CUDA
/* keep 3 GB VRAM headroom for the activation buffers + future tiers */
if (g_cuda && gpu_ok && ink_cuda_free_bytes() > (size_t)t->nbytes + (3ULL<<30)) {
w.dev = ink_cuda_upload(w.h, t->nbytes);
if (w.dev) { free(w.h); w.h = NULL; }
}
#else
(void)gpu_ok;
#endif
} else {
w.f = falloc(t->numel);
st_read_f32(&m->S, name, w.f, 0);
}
return w;
}
/* `off` conta ELEMENTI ed e' sempre un multiplo di I (si affetta per riga: un
* esperto condiviso dal tensore fuso [E,R,I]), quindi I arriva dal chiamante —
* st_tensor non porta le shape e non voglio indovinarle. */
static Wt wt_off_i(Wt w, int64_t off, int I) {
Wt r = w;
#ifdef COLI_VULKAN
r.vk_view = w.vk_view + off;
#endif
r.f = w.f ? w.f + off : NULL;
r.h = w.h ? w.h + off : NULL;
r.dev = w.dev ? (char*)w.dev + off*2 : NULL; /* dev is always bf16 */
if (w.q4) {
int64_t row = off / I;
if (w.qbits == 8) { r.q4 = w.q4 + off; r.qs = w.qs + row; r.qn = w.qn - off; }
else { r.q4 = w.q4 + off / 2; r.qs = w.qs + row * ((I + w.gs - 1) / w.gs);
r.qn = w.qn - off / 2; }
}
return r;
}
/* dequantizza UNA riga (int4-gs64: nibble +8, low = colonna pari, scala ogni gs) */