url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #242 from woolcoxm/feat/grouped-quant-fmt4

Browse Source 
Group-scaled int4 (fmt=4): one scale per 128 elements — fixes incoherent output
This commit is contained in:
Vincenzo
2026-07-15 08:25:29 +02:00
committed by GitHub
parent ff04b320d0 498ab0c20e
commit 3fd47b7bbd
2 changed files with 122 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/glm.c
+77 -7
View File
@@ -96,7 +96,7 @@ typedef struct {
* INT4 e' cio' che fa stare la densa residente nei 15 GB (0.5 byte/param). */
/* fmt: 0 F32, 1 INT8, 2 INT4 (2/byte), 3 INT2 (4/byte). q4 ospita sia int4 che int2 packed. */
typedef struct {
int fmt; float *qf; int8_t *q8; uint8_t *q4; float *s; int O, I;
int fmt; float *qf; int8_t *q8; uint8_t *q4; float *s; int O, I, gs; /* gs=group size (0=per-row, 128=grouped) */
#ifdef COLI_CUDA
ColiCudaTensor *cuda;
#endif
@@ -107,7 +107,10 @@ static int64_t qt_bytes(const QT *t){ /* byte residenti del tensore */
if(t->fmt==0) return n*4;
if(t->fmt==1) return n + (int64_t)t->O*4;
if(t->fmt==3) return (int64_t)t->O*((t->I+3)/4) + (int64_t)t->O*4;
return (int64_t)t->O*((t->I+1)/2) + (int64_t)t->O*4;
if(t->fmt==4){ /* int4 grouped: packed nibbles + O*ceil(I/gs) scales */
int ng=(t->I+t->gs-1)/t->gs;
return (int64_t)t->O*((t->I+1)/2) + (int64_t)t->O*ng*4; }
return (int64_t)t->O*((t->I+1)/2) + (int64_t)t->O*4; /* fmt=2 int4 per-row */
}
typedef struct {
@@ -380,6 +383,46 @@ static void matmul_i4(float *y, const float *x, const uint8_t *q4, const float *
if(i<I){ uint8_t byte=w[i>>1]; int lo=(int)(byte&0xF)-8; a += xs[i]*(float)lo; }
y[(int64_t)s*O+o]=a*sc; } }
}
/* y[S,O] = x[S,I] @ W^T with W int4 packed (2/byte) + per-GROUP scales (fmt=4).
* Same nibble math as matmul_i4, but the scale changes every `gs` elements along I.
* The accumulator resets at each group boundary: dot(x[grp], w[grp]) * scale[grp].
* gs MUST be a multiple of 16 (the AVX2 vector width). */
static void matmul_i4_grouped(float *y, const float *x, const uint8_t *q4, const float *scale,
int S, int I, int O, int gs){
int rb=(I+1)/2; int ng=(I+gs-1)/gs;
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){
const uint8_t *w=q4+(int64_t)o*rb;
const float *scl=scale+(int64_t)o*ng;
for(int s=0;s<S;s++){
const float *xs=x+(int64_t)s*I; float a=0;
for(int g=0; g*gs<I; g++){
int base=g*gs; int glen=gs; if(base+glen>I) glen=I-base;
float sc=scl[g];
int i=base;
#ifdef __AVX2__
const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8);
__m256 acc=_mm256_setzero_ps();
for(; i+16<=base+glen; i+=16){ __m128i by=_mm_loadl_epi64((const __m128i*)(w+(i>>1)));
__m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i nib=_mm_unpacklo_epi8(lo,hi);
__m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8));
__m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8));
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc);
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); }
a+=hsum256(acc)*sc;
#endif
/* scalar tail for the group remainder */
for(; i<base+glen; i+=2){
if(i+1<base+glen){ uint8_t byte=w[i>>1];
a+=(xs[i]*(float)((int)(byte&0xF)-8)+xs[i+1]*(float)((int)(byte>>4)-8))*sc; }
else { uint8_t byte=w[i>>1]; a+=xs[i]*(float)((int)(byte&0xF)-8)*sc; }
}
}
y[(int64_t)s*O+o]=a;
}
}
}
/* Decode hot path for gate+up: same exact q4 dot products as matmul_i4, but one
* OpenMP dispatch covers both matrices. KTransformers uses persistent pools;
* this keeps colibri dependency-free while removing one team launch/expert. */
@@ -772,6 +815,9 @@ static void matmul_qt_ex(float *y, const float *x, QT *w, int S, int allow_idot)
}
#endif
if(w->fmt==0){ matmul(y,x,w->qf,S,w->I,w->O); return; }
/* fmt=4: grouped int4 — always use the exact grouped kernel (no IDOT approximation,
* since the whole point of grouped scales is better quality). */
if(w->fmt==4){ matmul_i4_grouped(y,x,w->q4,w->s,S,w->I,w->O,w->gs); return; }
/* int8 IDOT vince sempre (1.4-2.5x). int4 IDOT: l'autore su AVX2 trovo' che a S=1
* non ripaga (soglia S>=2); ma su ARM/SDOT il singolo token CONVIENE (vedi g_i4s /
* PR #9 per il gemello VNNI). Soglia configurabile con I4S.
@@ -1072,9 +1118,22 @@ static void qt_from_disk(Model *m, const char *name, int O, int I, int bits, int
char sn[300]; snprintf(sn,sizeof(sn),"%s.qs",name);
if(st_has(&m->S,sn)){
int64_t nb=st_nbytes(&m->S,name);
int fmt = (nb==(int64_t)O*I)?1 : (nb==(int64_t)O*((I+1)/2))?2 : 3; /* int8 / int4 / int2 dai byte */
if(fmt==1){ if(t->fmt!=1||!t->q8){ t->fmt=1; t->O=O; t->I=I; t->q8=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q8,drop); }
else { if(t->fmt!=fmt||!t->q4){ t->fmt=fmt; t->O=O; t->I=I; t->q4=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q4,drop); }
int64_t ns=st_nbytes(&m->S,sn); /* scale bytes (F32) */
/* Detect int4-grouped (fmt=4): packed int4 weight bytes BUT scale array is
* larger than O*4 — check if it matches O*ceil(I/gs)*4 for gs=128. */
int fmt = (nb==(int64_t)O*I)?1 : (nb==(int64_t)O*((I+1)/2))?2 : 3;
int gs=0;
if(fmt==2 && ns > (int64_t)O*4){
/* could be grouped; try gs=128 */
int ng128=(I+127)/128;
if(ns==(int64_t)O*ng128*4){ fmt=4; gs=128; }
/* future: try other group sizes here */
}
if(fmt==1){ if(t->fmt!=1||!t->q8){ t->fmt=1; t->O=O; t->I=I; t->gs=0; t->q8=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q8,drop); }
else if(fmt==4){ int ng=(I+gs-1)/gs;
if(t->fmt!=4||!t->q4){ t->fmt=4; t->O=O; t->I=I; t->gs=gs; t->q4=qalloc(nb); t->s=falloc((int64_t)O*ng); }
st_read_raw(&m->S,name,t->q4,drop); }
else { if(t->fmt!=fmt||!t->q4){ t->fmt=fmt; t->O=O; t->I=I; t->gs=0; t->q4=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q4,drop); }
st_read_f32(&m->S,sn,t->s,drop);
} else {
if(!t->qf && !t->q8 && !t->q4) qt_alloc(t,O,I,bits);
@@ -1397,7 +1456,13 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal){
for(int k=0;k<3;k++){
int64_t nb=tw[k]->nbytes;
int fmt=(nb==(int64_t)OO[k]*II[k])?1:(nb==(int64_t)OO[k]*((II[k]+1)/2))?2:3;
qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->qf=NULL;
/* detect grouped int4 (fmt=4): int4 weight bytes + larger scale array */
int gs=0;
if(fmt==2 && tq[k]->nbytes > (int64_t)OO[k]*4){
int ng128=(II[k]+127)/128;
if(tq[k]->nbytes==(int64_t)OO[k]*ng128*4){ fmt=4; gs=128; }
}
qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->gs=gs; qt[k]->qf=NULL;
qt[k]->q8=(int8_t*)((char*)bw[k]+tw[k]->off); qt[k]->q4=(uint8_t*)((char*)bw[k]+tw[k]->off);
qt[k]->s=(float*)((char*)bq[k]+tq[k]->off);
}
@@ -1511,7 +1576,12 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal){
for(int k=0;k<3;k++){
int64_t nb=tw[k]->nbytes;
int fmt = (nb==(int64_t)OO[k]*II[k])?1 : (nb==(int64_t)OO[k]*((II[k]+1)/2))?2 : 3;
qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->qf=NULL;
int gs=0;
if(fmt==2 && tq[k]->nbytes > (int64_t)OO[k]*4){
int ng128=(II[k]+127)/128;
if(tq[k]->nbytes==(int64_t)OO[k]*ng128*4){ fmt=4; gs=128; }
}
qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->gs=gs; qt[k]->qf=NULL;
qt[k]->q8=(int8_t*)(s->slab+pos[k]); qt[k]->q4=s->slab+pos[k]; qt[k]->s=fp[k];
}
s->eid=eid; return 0;
c/tools/convert_fp8_to_int4.py
+45 -7
View File
@@ -51,6 +51,38 @@ def quant_int4(w, bits): # -> (qbytes U8 [O*ceil(I/2)], s
out[:, :v1.shape[1]] |= (v1 << 4)
return out.reshape(-1), s[:, 0].astype(np.float32)
def quant_int4_grouped(w, bits, gs=128):
"""Group-scaled int4: one scale per group of `gs` elements along the input dim.
Drastically reduces quantization error vs per-row scaling — matches the FP8
source's 128x128 block-scale granularity. Output layout:
qbytes: same packed nibble format as quant_int4
scales: f32 [O * ngroups] where ngroups = ceil(I/gs), laid out as
s[o * ngroups + g] = scale for row o, group g.
The engine detects this format (fmt=4) by checking the .qs array size."""
O, I = w.shape
qmax = (1 << (bits - 1)) - 1
ngroups = (I + gs - 1) // gs
# pad I to a multiple of gs for clean reshape, then trim
Ipad = ngroups * gs
wpad = np.zeros((O, Ipad), np.float32)
wpad[:, :I] = w
wr = wpad.reshape(O, ngroups, gs) # [O, ngroups, gs]
amax = np.abs(wr).max(axis=2, keepdims=True) # [O, ngroups, 1]
s = np.maximum(amax / qmax, 1e-8) # [O, ngroups, 1]
q = np.clip(np.rint(wr / s), -8, qmax).astype(np.int32) # [O, ngroups, gs]
q = q.reshape(O, Ipad)[:, :I] # trim padding -> [O, I]
# pack nibbles (identical to quant_int4)
rb = (I + 1) // 2
out = np.zeros((O, rb), np.uint8)
v0 = (q[:, 0::2] + 8).astype(np.uint8)
out[:, :v0.shape[1]] = v0
if I > 1:
v1 = (q[:, 1::2] + 8).astype(np.uint8)
out[:, :v1.shape[1]] |= (v1 << 4)
# scales: flatten [O, ngroups] -> [O * ngroups]
s_flat = s[:, :, 0].astype(np.float32).reshape(-1)
return out.reshape(-1), s_flat
def quant_int2(w, bits): # -> (qbytes U8 [O*ceil(I/4)], scale f32 [O]); 4/byte
O, I = w.shape
qmax = (1 << (bits - 1)) - 1 # bits=2 -> qmax=1, valori [-2,1]
@@ -169,7 +201,8 @@ def dequant(f, name, keys):
return (w * sc).numpy()
return f.get_tensor(name).to(torch.float32).numpy()
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits, keep_mtp=False, keep_idx=False):
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits,
keep_mtp=False, keep_idx=False, group_size=0):
from safetensors import safe_open
with safe_open(path, framework="pt") as f:
keys = set(f.keys())
@@ -183,8 +216,11 @@ def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits, keep_mtp=Fals
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
if w.ndim != 2: # es. bias 1D non previsto come 'q' -> tienilo f32
out_dict[name] = w.astype(np.float32); continue
q, s = (quant_int2(w, bits) if bits <= 2 else
quant_int4(w, bits) if bits <= 4 else quant_int8(w, bits))
if group_size > 0 and bits <= 4:
q, s = quant_int4_grouped(w, bits, group_size)
else:
q, s = (quant_int2(w, bits) if bits <= 2 else
quant_int4(w, bits) if bits <= 4 else quant_int8(w, bits))
out_dict[name] = q
out_dict[name + ".qs"] = s
@@ -198,6 +234,8 @@ def main():
ap.add_argument("--ebits", type=int, default=None) # bit residenti (default 4; 8 per --mtp/--indexer)
ap.add_argument("--io-bits", type=int, default=8) # bit di embed/lm_head
ap.add_argument("--xbits", type=int, default=None) # bit degli expert ROUTED (streaming); default=ebits
ap.add_argument("--group-size", type=int, default=0, # 0 = per-row (backward compat); 128 = group-scaled
help="group size for int4 scales: 0=per-row (default), 128=one scale per 128 elements (much better quality)")
ap.add_argument("--n-layers", type=int, default=78)
ap.add_argument("--min-free-gb", type=float, default=20.0)
ap.add_argument("--selftest", action="store_true")
@@ -298,7 +336,7 @@ def main():
shards = sorted(glob.glob(os.path.join(a.indir, "*.safetensors")))
from safetensors.numpy import save_file
for i, sp in enumerate(shards):
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits)
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size)
save_file(out, os.path.join(a.outdir, f"out-{i:05d}.safetensors"))
# copia config + tokenizer
for fn in ["config.json"]:
@@ -541,7 +579,7 @@ def main():
if os.path.exists(outp): print(f"[MTP] {outp} already done"); continue
print(f"[MTP {i+1}/{len(mtp_shards)}] downloading {sh}...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True, group_size=a.group_size)
save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -561,7 +599,7 @@ def main():
if os.path.exists(outp): continue # gia' fatto -> ripartibile
print(f"[IDX {i+1}/{len(idx_shards)}] downloading {sh}...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True, group_size=a.group_size)
if out: save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -575,7 +613,7 @@ def main():
if os.path.exists(outp): continue # gia' fatto -> ripartibile
print(f"[{i+1}/{len(shards)}] downloading {sh} ({free_gb(a.outdir):.0f} GB free)...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size)
save_file(out, outp)
os.remove(p) # <-- cancella subito lo shard fp8
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):