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

diskio: KV write batching + persistent handle; generators: unfuse experts

Browse Source 
Two independent fixes validated end-to-end on fresh fixtures:

1. KV cache disk I/O (issue_diskio.md opportunities #1 + #4):
   - kv_disk_append: fopen/fclose every turn -> persistent FILE* kept open
     for the engine lifetime, lazy open on first append, closed in
     serve_ctx_free. Eliminates per-turn handle creation overhead.
   - kv_disk_append: ~157 small fwrites per position -> one contiguous
     record memcpy'd into a staging buffer then a single fwrite per
     position. The staging buffer grows on demand via realloc.
   - kv_disk_truncate: closes the persistent handle before truncating
     so the file actually shrinks on disc, then reopens lazily.
   - KVState gains disk_fp, disk_buf, disk_buf_cap fields.
   - Verified: serve-mode round-trip, write 11 tokens then reload and
     resume with no re-prefill, then append 8 more and reload to 19.

2. Expert weight unfusing in test-model generators:
   - The real GLM-5.2-FP8 checkpoint stores routed experts UNFUSED as
     per-expert 2-D tensors, each with its own _scale_inv. HF fuses
     gate+up into a single 3-D gate_up_proj for compute efficiency.
   - The converter and C engine both expect the unfused layout. The
     fused 3-D tensors were silently skipped by the converter, and the
     engine crashed with missing-tensor errors.
   - New unfuse_experts in glm_fp8_emit.py splits gate_up_proj and
     down_proj into per-expert 2-D tensors. Called after reference
     generation but before saving, in both generators, both FP8 and bf16.
   - Also fixed: make_glm_oracle.py FP8 round-trip guard used p.dim()<2
     which let 3-D fused experts through and crashed fp8_block_quantize.
     Changed to p.dim()!=2 to match the converter ndim!=2 guard.

Validated full chain on fresh fixtures:
  generator --fp8 -> 570 e4m3 tensors + 629 scale_inv, was 90 when fused
  converter --group-size 0  -> per-row int4 fmt=2, engine loads clean
  converter --group-size 128 -> grouped int4 fmt=4, 8-16x more scales,
    engine loads clean, fmt=4 auto-detected in both mmap and slab paths
  dequant error: grouped 1.14-1.22x lower than per-row vs FP8 source
This commit is contained in:
woolcoxm
2026-07-15 01:40:02 -04:00
parent e71d4fbe29
commit 71c262ce1a
4 changed files with 127 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/glm.c
+54 -12
View File
@@ -142,6 +142,9 @@ typedef struct {
int *kv_start, max_t;
int disk_nrec;
char disk_path[2048];
FILE *disk_fp; /* kept-open handle: fopen once, fwrite per turn, fclose at exit (#4) */
uint8_t *disk_buf; /* staging buffer: one contiguous record per position (#1) */
int64_t disk_buf_cap;
} KVState;
typedef struct {
@@ -3971,36 +3974,73 @@ static void kv_hdr(Model *m, int32_t *h, int nrec){
h[0]=c->n_layers; h[1]=c->kv_lora; h[2]=c->qk_rope;
h[3]=m->has_dsa?c->index_hd:0; h[4]=nic; h[5]=c->vocab; h[6]=nrec; h[7]=0;
}
/* Bytes of one on-disk record: [tok i32][Lc+Rc per layer][Ic per DSA layer].
* Layout matches what kv_disk_append writes and kv_disk_load reads. */
static int64_t kv_rec_bytes(Model *m){
Cfg *c=&m->c;
int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4;
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4;
return rec;
}
/* Open the persistent handle lazily; write the header if the file is new. After
* this returns successfully, k->disk_fp is valid for the engine's lifetime and
* positioned at end-of-header (nrec==0 case) or wherever the caller seeks. */
static int kv_disk_open(Model *m){
KVState *k=m->kv;
if(k->disk_fp) return 1;
k->disk_fp=fopen(k->disk_path,"r+b");
if(!k->disk_fp){ /* not there yet -> create + header */
k->disk_fp=fopen(k->disk_path,"wb");
if(!k->disk_fp) return 0;
int32_t h[8]; kv_hdr(m,h,0);
fwrite(KV_MAGIC,1,8,k->disk_fp); fwrite(h,4,8,k->disk_fp);
fflush(k->disk_fp);
fclose(k->disk_fp);
k->disk_fp=fopen(k->disk_path,"r+b"); /* reopen r+b for append */
if(!k->disk_fp) return 0;
}
return 1;
}
static void kv_disk_truncate(Model *m, int nrec){
if(!g_kvsave) return;
KVState *k=m->kv;
if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; } /* drop to shrink on disc */
FILE *f=fopen(k->disk_path,"r+b");
if(!f){ k->disk_nrec=0; return; }
k->disk_nrec=nrec;
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f);
fflush(f); fclose(f);
}
static void kv_disk_reset(Model *m){ kv_disk_truncate(m,0); }
static void kv_disk_append(Model *m, const int *hist, int len){
KVState *k=m->kv;
if(!g_kvsave || len<=k->disk_nrec) return;
Cfg *c=&m->c;
FILE *f=fopen(k->disk_path,"r+b");
if(!f){ f=fopen(k->disk_path,"wb"); if(!f) return;
int32_t h[8]; kv_hdr(m,h,0); fwrite(KV_MAGIC,1,8,f); fwrite(h,4,8,f); }
int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4;
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4;
if(!kv_disk_open(m)) return;
FILE *f=k->disk_fp;
int64_t rec = kv_rec_bytes(m);
/* grow the contiguous staging buffer if the record is larger (#1 batching) */
if(rec > k->disk_buf_cap){
uint8_t *nb=realloc(k->disk_buf, rec);
if(!nb) return; /* OOM: skip this turn, retry next */
k->disk_buf=nb; k->disk_buf_cap=rec;
}
fseek(f, 8+8*4 + (int64_t)k->disk_nrec*rec, SEEK_SET);
for(int p=k->disk_nrec;p<len;p++){
int32_t tk=hist[p]; fwrite(&tk,4,1,f);
uint8_t *b=k->disk_buf; /* pack token + every layer into one record */
*(int32_t*)b = hist[p]; b+=4;
for(int i=0;i<c->n_layers;i++){
fwrite(m->Lc[i]+(int64_t)p*c->kv_lora, 4, c->kv_lora, f);
fwrite(m->Rc[i]+(int64_t)p*c->qk_rope, 4, c->qk_rope, f);
memcpy(b, m->Lc[i]+(int64_t)p*c->kv_lora, (size_t)c->kv_lora*4); b+=c->kv_lora*4;
memcpy(b, m->Rc[i]+(int64_t)p*c->qk_rope,(size_t)c->qk_rope*4); b+=c->qk_rope*4;
}
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i])
fwrite(m->Ic[i]+(int64_t)p*c->index_hd, 4, c->index_hd, f);
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]){
memcpy(b, m->Ic[i]+(int64_t)p*c->index_hd, (size_t)c->index_hd*4); b+=c->index_hd*4;
}
fwrite(k->disk_buf, 1, (size_t)rec, f); /* one fwrite per position (was ~157) */
}
fflush(f); /* dati prima, contatore poi */
int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f);
fflush(f); /* persist the counter too */
k->disk_nrec=len;
}
static int kv_disk_load(Model *m, int *hist, int maxctx){
@@ -4055,6 +4095,8 @@ static void serve_ctx_init(Model *m, ServeCtx *s, const char *snap, int slot, in
static void serve_ctx_free(Model *m, ServeCtx *s){
KVState *k=&s->kv; int NR=m->c.n_layers+1;
if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; }
free(k->disk_buf); k->disk_buf=NULL;
if(k->Lc) for(int i=0;i<NR;i++){ free(k->Lc[i]); free(k->Rc[i]); }
if(k->Ic) for(int i=0;i<m->c.n_layers;i++) free(k->Ic[i]);
free(k->Lc); free(k->Rc); free(k->Ic); free(k->kv_start); free(s->hist);
c/tools/glm_fp8_emit.py
+42
View File
@@ -66,6 +66,48 @@ def fp8_block_dequantize(w_fp8, scale):
return qf * scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)[:O, :I]
def unfuse_experts(sd):
"""Split HF's fused 3-D `experts.gate_up_proj` [E, 2*M, I] into per-expert 2-D
`experts.{e}.gate_proj` [M, I] + `experts.{e}.up_proj` [M, I], and
`experts.down_proj` [E, I, M] -> `experts.{e}.down_proj` [M_out, I].
The real GLM-5.2-FP8 checkpoint stores experts UNFUSED as per-expert 2-D tensors
(gate_proj, up_proj, down_proj), each with its own _scale_inv. HF's
GlmMoeDsaForCausalLM fuses gate+up into a single 3-D gate_up_proj for efficiency.
The converter (classify + ndim!=2 guard) and the C engine both expect the unfused
layout, so we split before saving.
Idempotent: if experts are already unfused (no 3-D gate_up_proj), returns sd as-is.
EN: split HF's fused 3-D expert weights into the per-expert 2-D layout that the real
EN: checkpoint uses and the converter/engine expect. No-op if already unfused."""
keys_to_remove = []
new_entries = {}
for name, t in sd.items():
if not name.endswith(".mlp.experts.gate_up_proj"):
continue
# prefix = everything before ".mlp.experts.gate_up_proj"
prefix = name[:-len(".mlp.experts.gate_up_proj")]
E, twoM, I = t.shape # [E, 2*intermediate, input]
M = twoM // 2
for e in range(E):
new_entries[f"{prefix}.mlp.experts.{e}.gate_proj.weight"] = t[e, :M, :].contiguous()
new_entries[f"{prefix}.mlp.experts.{e}.up_proj.weight"] = t[e, M:, :].contiguous()
keys_to_remove.append(name)
# down_proj may be 3-D [E, I, M] in the fused form, or already per-expert
for name, t in sd.items():
if not name.endswith(".mlp.experts.down_proj") or t.dim() != 3:
continue
prefix = name[:-len(".mlp.experts.down_proj")]
E = t.shape[0]
for e in range(E):
new_entries[f"{prefix}.mlp.experts.{e}.down_proj.weight"] = t[e].contiguous()
keys_to_remove.append(name)
for k in keys_to_remove:
sd.pop(k, None)
sd.update(new_entries)
return sd
def state_dict_to_fp8(sd):
"""Converte uno state_dict HuggingFace nel layout FP8 del checkpoint reale:
per ogni tensore quantizzabile 2-D scrive `{name}` (F8_E4M3) + `{name}_scale_inv` (F32);
c/tools/make_glm_bench_model.py
+20 -11
View File
@@ -22,7 +22,7 @@ import torch
from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/
from glm_fp8_emit import save_fp8_safetensors
from glm_fp8_emit import save_fp8_safetensors, unfuse_experts
def build_config() -> GlmMoeDsaConfig:
@@ -85,16 +85,6 @@ def main() -> None:
output = Path(args.output)
output.mkdir(parents=True, exist_ok=True)
params = sum(p.numel() for p in model.parameters())
if args.fp8:
n_fp8, n_tot = save_fp8_safetensors(model.state_dict(), output / "model.safetensors")
# save_pretrained scrive config.json; nel path FP8 lo bypassiamo, quindi lo scriviamo
# a mano (serve al converter e al motore C). EN: save_pretrained writes config.json;
# the FP8 path bypasses it, so write it manually (converter + C engine need it).
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
print(f"saved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
f"-> {output / 'model.safetensors'}")
else:
model.save_pretrained(output, safe_serialization=True, max_shard_size="4GB")
model.to(args.device)
prompt = [3, 14, 159, 26, 53, 58, 200, 11, 77, 240, 5, 99]
@@ -103,6 +93,25 @@ def main() -> None:
full = model.generate(ids, max_new_tokens=8, do_sample=False, use_cache=True)[0]
logits = model(full.unsqueeze(0), use_cache=False).logits[0]
# Unfuse experts AFTER reference generation (model needs fused weights for
# forward/generate) but BEFORE saving — the real checkpoint and the converter
# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors.
sd = model.state_dict()
unfuse_experts(sd)
if args.fp8:
n_fp8, n_tot = save_fp8_safetensors(sd, output / "model.safetensors")
# save_pretrained scrive config.json; nel path FP8 lo bypassiamo, quindi lo scriviamo
# a mano (serve al converter e al motore C). EN: save_pretrained writes config.json;
# the FP8 path bypasses it, so write it manually (converter + C engine need it).
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
print(f"saved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
f"-> {output / 'model.safetensors'}")
else:
from safetensors.torch import save_file
save_file({k: v.contiguous() for k, v in sd.items()}, str(output / "model.safetensors"))
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
ref = {
"prompt_ids": prompt,
"full_ids": full.cpu().tolist(),
c/tools/make_glm_oracle.py
+11 -4
View File
@@ -22,7 +22,7 @@ from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/
from glm_fp8_emit import (fp8_block_quantize, fp8_block_dequantize, keep_f32,
save_fp8_safetensors)
save_fp8_safetensors, unfuse_experts)
ap = argparse.ArgumentParser()
ap.add_argument("--fp8", action="store_true",
@@ -84,7 +84,7 @@ with torch.no_grad():
if args.fp8:
with torch.no_grad():
for n, p in model.named_parameters():
if keep_f32(n, p) or p.dim() < 2:
if keep_f32(n, p) or p.dim() != 2:
continue
q, s = fp8_block_quantize(p)
p.copy_(fp8_block_dequantize(q, s))
@@ -109,12 +109,19 @@ with torch.no_grad():
tf_pred = lg.argmax(-1).tolist()
print("tf_pred:", tf_pred)
# Unfuse experts AFTER reference generation (model needs fused weights for
# forward/generate) but BEFORE saving — the real checkpoint and the converter
# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors.
sd = model.state_dict()
unfuse_experts(sd)
if args.fp8:
n_fp8, n_tot = save_fp8_safetensors(model.state_dict(), "glm_tiny/model.safetensors")
n_fp8, n_tot = save_fp8_safetensors(sd, "glm_tiny/model.safetensors")
print(f"\nsaved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
f"-> glm_tiny/model.safetensors")
else:
model.save_pretrained("glm_tiny", safe_serialization=True)
from safetensors.torch import save_file
save_file({k: v.contiguous() for k, v in sd.items()}, "glm_tiny/model.safetensors")
json.dump(cfg.to_dict(), open("glm_tiny/config.json", "w"))
json.dump({"prompt_ids": prompt, "full_ids": full, "tf_pred": tf_pred}, open("ref_glm.json", "w"))
print("saved: glm_tiny/ (weights + config) and ref_glm.json"