#!/usr/bin/env python3
"""
Streaming GGUF expert slicer for Qwen3.8-Flash-Next (qwen4exp).
Slices blk.{l}.ffn_{down,gate,up}_exps.weight and blk.{l}.ffn_gate_inp.weight
from N_experts (512) down to K experts per layer using a per-layer keep list,
updates qwen4exp.expert_count = K, and streams non-expert tensors via copy_file_range/chunks
with zero re-quantization and no enum dependency on gguf-py.
"""
import argparse
import json
import os
import re
import struct
import sys
import tempfile
from pathlib import Path
GGUF_MAGIC = 0x46554747 # "GGUF" little-endian
DEFAULT_ALIGNMENT = 32
# GGUF metadata value types
GGUF_TYPE_UINT8 = 0
GGUF_TYPE_INT8 = 1
GGUF_TYPE_UINT16 = 2
GGUF_TYPE_INT16 = 3
GGUF_TYPE_UINT32 = 4
GGUF_TYPE_INT32 = 5
GGUF_TYPE_FLOAT32 = 6
GGUF_TYPE_BOOL = 7
GGUF_TYPE_STRING = 8
GGUF_TYPE_ARRAY = 9
GGUF_TYPE_UINT64 = 10
GGUF_TYPE_INT64 = 11
GGUF_TYPE_FLOAT64 = 12
SCALAR_SIZES = {
GGUF_TYPE_UINT8: 1,
GGUF_TYPE_INT8: 1,
GGUF_TYPE_UINT16: 2,
GGUF_TYPE_INT16: 2,
GGUF_TYPE_UINT32: 4,
GGUF_TYPE_INT32: 4,
GGUF_TYPE_FLOAT32: 4,
GGUF_TYPE_BOOL: 1,
GGUF_TYPE_UINT64: 8,
GGUF_TYPE_INT64: 8,
GGUF_TYPE_FLOAT64: 8,
}
EXPERT_TENSOR_RE = re.compile(
r"^blk\.(\d+)\.(ffn_(?:down|gate|up|gate_up)_exps\.(?:weight|bias|scale)|ffn_gate_inp\.(?:weight|bias)|exp_probs_b\.bias)$"
)
NON_EXPERT_512_SUFFIXES = (".attn_k.weight", ".attn_v.weight", ".indexer.q_proj.weight")
def pwrite_all(fd: int, data: bytes | memoryview, offset: int) -> None:
mv = memoryview(data)
written = 0
total = len(mv)
while written < total:
n = os.pwrite(fd, mv[written:], offset + written)
if n <= 0:
raise OSError(f"pwrite returned {n} at offset {offset + written}")
written += n
def align_offset(offset: int, alignment: int = DEFAULT_ALIGNMENT) -> int:
rem = offset % alignment
return offset if rem == 0 else offset + (alignment - rem)
class BinaryCursor:
def __init__(self, fd: int):
self.fd = fd
self.pos = 0
self.buf = b""
self.buf_start = 0
def _ensure(self, n: int) -> bytes:
if self.pos + n > self.buf_start + len(self.buf):
read_len = max(n, 256 * 1024)
chunk = os.pread(self.fd, read_len, self.pos)
if len(chunk) < n:
raise EOFError(f"Unexpected EOF at {self.pos}, wanted {n} bytes")
self.buf = chunk
self.buf_start = self.pos
off = self.pos - self.buf_start
res = self.buf[off : off + n]
self.pos += n
return res
def read_u32(self) -> int:
return struct.unpack("<I", self._ensure(4))[0]
def read_u64(self) -> int:
return struct.unpack("<Q", self._ensure(8))[0]
def read_string_raw(self) -> tuple[str, bytes]:
raw_len = self._ensure(8)
slen = struct.unpack("<Q", raw_len)[0]
raw_bytes = self._ensure(slen)
return raw_bytes.decode("utf-8", errors="replace"), raw_len + raw_bytes
def read_value_raw(self, vtype: int) -> bytes:
if vtype in SCALAR_SIZES:
return self._ensure(SCALAR_SIZES[vtype])
if vtype == GGUF_TYPE_STRING:
_, raw = self.read_string_raw()
return raw
if vtype == GGUF_TYPE_ARRAY:
hdr = self._ensure(12)
elem_type, count = struct.unpack("<IQ", hdr)
if elem_type in SCALAR_SIZES:
body = self._ensure(SCALAR_SIZES[elem_type] * count)
return hdr + body
parts = [hdr]
for _ in range(count):
parts.append(self.read_value_raw(elem_type))
return b"".join(parts)
raise ValueError(f"Unknown GGUF metadata value type: {vtype}")
def parse_gguf_header(fd: int, file_size: int) -> dict:
cur = BinaryCursor(fd)
magic = cur.read_u32()
if magic != GGUF_MAGIC:
raise ValueError(f"Not a GGUF file (magic={hex(magic)})")
version = cur.read_u32()
n_tensors = cur.read_u64()
n_kv = cur.read_u64()
alignment = DEFAULT_ALIGNMENT
expert_count = None
kvs = []
for _ in range(n_kv):
key, raw_key = cur.read_string_raw()
raw_vtype = cur._ensure(4)
vtype = struct.unpack("<I", raw_vtype)[0]
raw_val = cur.read_value_raw(vtype)
if key == "general.alignment" and vtype == GGUF_TYPE_UINT32:
alignment = struct.unpack("<I", raw_val)[0]
elif key.endswith(".expert_count") and vtype == GGUF_TYPE_UINT32:
expert_count = struct.unpack("<I", raw_val)[0]
kvs.append({
"key": key,
"raw_key": raw_key,
"vtype": vtype,
"raw_vtype": raw_vtype,
"raw_val": raw_val,
})
tensors = []
for _ in range(n_tensors):
name, raw_name = cur.read_string_raw()
n_dims = cur.read_u32()
dims = [cur.read_u64() for _ in range(n_dims)]
ggml_type = cur.read_u32()
rel_offset = cur.read_u64()
tensors.append({
"name": name,
"raw_name": raw_name,
"n_dims": n_dims,
"dims": dims,
"ggml_type": ggml_type,
"rel_offset": rel_offset,
})
data_section_offset = align_offset(cur.pos, alignment)
# Determine exact byte size of each tensor from sorted offsets and file size
# Note: tensors in GGUF are laid out in ascending rel_offset order, each padded to alignment
for i, t in enumerate(tensors):
abs_start = data_section_offset + t["rel_offset"]
t["abs_start"] = abs_start
if i + 1 < len(tensors):
next_start = data_section_offset + tensors[i + 1]["rel_offset"]
# Tensor size is next_start minus alignment padding; for expert tensors we compute exact slice size
t["padded_span"] = next_start - abs_start
else:
t["padded_span"] = file_size - abs_start
return {
"version": version,
"n_tensors": n_tensors,
"n_kv": n_kv,
"alignment": alignment,
"expert_count": expert_count,
"kvs": kvs,
"tensors": tensors,
"data_section_offset": data_section_offset,
}
# Exact block size (in elements) and byte size per block for all GGML types used in GGUFs
GGML_TYPE_TRAITS = {
0: (1, 4), # F32
1: (1, 2), # F16
2: (32, 18), # Q4_0
3: (32, 20), # Q4_1
6: (32, 22), # Q5_0
7: (32, 24), # Q5_1
8: (32, 34), # Q8_0
9: (32, 36), # Q8_1
10: (256, 84), # Q2_K
11: (256, 110), # Q3_K
12: (256, 144), # Q4_K
13: (256, 176), # Q5_K
14: (256, 210), # Q6_K
15: (256, 292), # Q8_K
16: (256, 66), # IQ2_XXS
17: (256, 74), # IQ2_XS
18: (256, 98), # IQ3_XXS
19: (256, 50), # IQ1_S
20: (32, 18), # IQ4_NL
21: (256, 110), # IQ3_S
22: (256, 82), # IQ2_S
23: (256, 136), # IQ4_XS
24: (1, 1), # I8
25: (1, 2), # I16
26: (1, 4), # I32
27: (1, 8), # I64
28: (1, 8), # F64
29: (256, 56), # IQ1_M
30: (1, 2), # BF16
34: (256, 54), # TQ1_0
35: (256, 66), # TQ2_0
39: (32, 17), # MXFP4
40: (64, 36), # NVFP4
41: (128, 18), # Q1_0
42: (64, 18), # Q2_0 (64 weights = 16B qs + 2B fp16 d = 18B, 2.25 bpw)
}
def compute_tensor_nbytes(dims: list[int], ggml_type: int) -> int:
if ggml_type not in GGML_TYPE_TRAITS:
raise ValueError(f"Unsupported ggml_type={ggml_type}")
blck_size, type_size = GGML_TYPE_TRAITS[ggml_type]
n_elems = 1
for d in dims:
n_elems *= d
if n_elems % blck_size != 0:
raise ValueError(f"Tensor dims {dims} not divisible by block size {blck_size} for type {ggml_type}")
return (n_elems // blck_size) * type_size
def copy_range_fast(src_fd: int, src_off: int, dst_fd: int, dst_off: int, length: int) -> None:
"""Fast streaming copy using copy_file_range on Linux or 64 MiB pread/pwrite chunks."""
has_cfr = hasattr(os, "copy_file_range")
rem = length
cur_s = src_off
cur_d = dst_off
chunk_sz = 64 * 1024 * 1024
while rem > 0:
step = min(rem, chunk_sz)
if has_cfr:
try:
n = os.copy_file_range(src_fd, dst_fd, step, offset_src=cur_s, offset_dst=cur_d)
if n > 0:
cur_s += n
cur_d += n
rem -= n
continue
has_cfr = False
except OSError:
has_cfr = False
buf = os.pread(src_fd, step, cur_s)
if not buf:
raise EOFError(f"Unexpected EOF copying {step} bytes at {cur_s}")
pwrite_all(dst_fd, buf, cur_d)
cur_s += len(buf)
cur_d += len(buf)
rem -= len(buf)
def slice_gguf_shard(
in_path: str,
out_path: str,
keep_map: dict[int, list[int]],
old_n_experts: int = 512,
) -> dict:
if os.path.abspath(in_path) == os.path.abspath(out_path):
raise ValueError("in_path and out_path must be distinct files")
file_size = os.path.getsize(in_path)
in_fd = os.open(in_path, os.O_RDONLY)
try:
hdr = parse_gguf_header(in_fd, file_size)
alignment = hdr["alignment"]
# Validate keep_map has uniform K across all present layers
ks = {len(v) for v in keep_map.values()}
if len(ks) != 1:
raise ValueError(f"GGUF requires uniform K across layers, got lengths {ks}")
new_k = next(iter(ks))
for l, idx_list in keep_map.items():
if len(set(idx_list)) != new_k or any(i < 0 or i >= old_n_experts for i in idx_list):
raise ValueError(f"Invalid expert indices at layer {l}")
# Build updated KV section
kv_bytes_list = []
for kv in hdr["kvs"]:
raw_val = kv["raw_val"]
if kv["key"].endswith(".expert_used_count") and kv["vtype"] == GGUF_TYPE_UINT32:
used_k = struct.unpack("<I", raw_val)[0]
if used_k > new_k:
raise ValueError(f"{kv['key']}={used_k} exceeds new expert count K={new_k}")
if kv["key"].endswith(".expert_count") and kv["vtype"] == GGUF_TYPE_UINT32:
raw_val = struct.pack("<I", new_k)
kv_bytes_list.append(kv["raw_key"] + kv["raw_vtype"] + raw_val)
# Build updated tensor infos and compute new relative offsets
new_tensors = []
cur_rel_offset = 0
sliced_count = 0
bytes_saved = 0
for t in hdr["tensors"]:
dims = list(t["dims"])
orig_nbytes = compute_tensor_nbytes(dims, t["ggml_type"])
m = EXPERT_TENSOR_RE.match(t["name"])
is_expert_tensor = False
layer_idx = None
if m is not None and dims[-1] == old_n_experts:
layer_idx = int(m.group(1))
if layer_idx not in keep_map:
raise KeyError(f"Layer {layer_idx} missing from keep_map for tensor {t['name']}")
dims[-1] = new_k
is_expert_tensor = True
sliced_count += 1
elif t["name"].startswith("blk.") and dims[-1] == old_n_experts and not t["name"].endswith(NON_EXPERT_512_SUFFIXES):
raise ValueError(f"Unhandled blk tensor with outermost dim {old_n_experts}: {t['name']} shape={dims}")
new_nbytes = compute_tensor_nbytes(dims, t["ggml_type"])
bytes_saved += (orig_nbytes - new_nbytes)
new_tensors.append({
"name": t["name"],
"raw_name": t["raw_name"],
"n_dims": t["n_dims"],
"orig_dims": t["dims"],
"new_dims": dims,
"ggml_type": t["ggml_type"],
"orig_abs_start": t["abs_start"],
"orig_nbytes": orig_nbytes,
"new_rel_offset": cur_rel_offset,
"new_nbytes": new_nbytes,
"is_expert_tensor": is_expert_tensor,
"layer_idx": layer_idx,
})
cur_rel_offset = align_offset(cur_rel_offset + new_nbytes, alignment)
# Serialize header + KV + tensor directory
out_hdr_parts = [
struct.pack("<IIQQ", GGUF_MAGIC, hdr["version"], hdr["n_tensors"], hdr["n_kv"]),
*kv_bytes_list,
]
for nt in new_tensors:
dims_packed = struct.pack(f"<{nt['n_dims']}Q", *nt["new_dims"])
info_packed = (
nt["raw_name"]
+ struct.pack("<I", nt["n_dims"])
+ dims_packed
+ struct.pack("<IQ", nt["ggml_type"], nt["new_rel_offset"])
)
out_hdr_parts.append(info_packed)
raw_header = b"".join(out_hdr_parts)
new_data_offset = align_offset(len(raw_header), alignment)
pad_len = new_data_offset - len(raw_header)
if pad_len > 0:
raw_header += b"\x00" * pad_len
out_fd = os.open(out_path, os.O_CREAT | os.O_TRUNC | os.O_WRONLY, 0o644)
try:
pwrite_all(out_fd, raw_header, 0)
for nt in new_tensors:
dst_start = new_data_offset + nt["new_rel_offset"]
if not nt["is_expert_tensor"]:
copy_range_fast(in_fd, nt["orig_abs_start"], out_fd, dst_start, nt["orig_nbytes"])
else:
# Gather K kept expert slices via per-expert pread (safe for tensors >2 GiB)
per_exp_bytes = nt["orig_nbytes"] // old_n_experts
keep_idxs = keep_map[nt["layer_idx"]]
sliced_chunks = []
for idx in keep_idxs:
exp_off = nt["orig_abs_start"] + idx * per_exp_bytes
chunk = os.pread(in_fd, per_exp_bytes, exp_off)
if len(chunk) != per_exp_bytes:
raise EOFError(f"Short read on {nt['name']} expert {idx}")
sliced_chunks.append(chunk)
out_bytes = b"".join(sliced_chunks)
assert len(out_bytes) == nt["new_nbytes"]
pwrite_all(out_fd, out_bytes, dst_start)
# Zero-pad to alignment boundary
aligned_end = new_data_offset + align_offset(nt["new_rel_offset"] + nt["new_nbytes"], alignment)
actual_end = dst_start + nt["new_nbytes"]
if aligned_end > actual_end:
pwrite_all(out_fd, b"\x00" * (aligned_end - actual_end), actual_end)
finally:
os.close(out_fd)
return {
"in_path": in_path,
"out_path": out_path,
"old_n_experts": old_n_experts,
"new_n_experts": new_k,
"sliced_tensors": sliced_count,
"bytes_saved": bytes_saved,
"out_file_size": os.path.getsize(out_path),
}
finally:
os.close(in_fd)
def _pack_str(s: str) -> bytes:
b = s.encode("utf-8")
return struct.pack("<Q", len(b)) + b
def run_self_test() -> None:
"""Create a synthetic qwen4exp GGUF with Q2_0 (type 42) 3D expert tensors, slice 8 -> 4 experts, verify."""
with tempfile.TemporaryDirectory() as td:
src = os.path.join(td, "test_in.gguf")
dst = os.path.join(td, "test_out.gguf")
old_e = 8
new_e = 4
keep = {0: [1, 3, 4, 7]}
# Build synthetic GGUF:
# KV 1: general.architecture = "qwen4exp"
# KV 2: qwen4exp.expert_count = 8 (uint32)
kv1 = _pack_str("general.architecture") + struct.pack("<I", GGUF_TYPE_STRING) + _pack_str("qwen4exp")
kv2 = _pack_str("qwen4exp.expert_count") + struct.pack("<I", GGUF_TYPE_UINT32) + struct.pack("<I", old_e)
# Tensor 0: token_embd.weight [64, 16] F32
t0_dims = [64, 16]
t0_bytes = os.urandom(compute_tensor_nbytes(t0_dims, 0))
# Tensor 1: blk.0.ffn_down_exps.weight [64, 32, 8] Q2_0 (type 42: 32 elems = 10B)
t1_dims = [64, 32, old_e]
t1_bytes_per_exp = (64 * 32 // 64) * 18
t1_exp_slices = [os.urandom(t1_bytes_per_exp) for _ in range(old_e)]
t1_bytes = b"".join(t1_exp_slices)
# Tensor 2: blk.0.ffn_gate_inp.weight [64, 8] F32
t2_dims = [64, old_e]
t2_bytes_per_exp = 64 * 4
t2_exp_slices = [os.urandom(t2_bytes_per_exp) for _ in range(old_e)]
t2_bytes = b"".join(t2_exp_slices)
off0 = 0
off1 = align_offset(off0 + len(t0_bytes))
off2 = align_offset(off1 + len(t1_bytes))
ti0 = _pack_str("token_embd.weight") + struct.pack("<I2QIQ", 2, *t0_dims, 0, off0)
ti1 = _pack_str("blk.0.ffn_down_exps.weight") + struct.pack("<I3QIQ", 3, *t1_dims, 42, off1)
ti2 = _pack_str("blk.0.ffn_gate_inp.weight") + struct.pack("<I2QIQ", 2, *t2_dims, 0, off2)
hdr = struct.pack("<IIQQ", GGUF_MAGIC, 3, 3, 2) + kv1 + kv2 + ti0 + ti1 + ti2
data_off = align_offset(len(hdr))
with open(src, "wb") as f:
f.write(hdr + b"\x00" * (data_off - len(hdr)))
f.seek(data_off + off0)
f.write(t0_bytes)
f.seek(data_off + off1)
f.write(t1_bytes)
f.seek(data_off + off2)
f.write(t2_bytes)
res = slice_gguf_shard(src, dst, keep, old_n_experts=old_e)
assert res["sliced_tensors"] == 2, res
out_fd = os.open(dst, os.O_RDONLY)
try:
parsed = parse_gguf_header(out_fd, os.path.getsize(dst))
assert parsed["expert_count"] == new_e
doff = parsed["data_section_offset"]
# Check t0 bit-identical
got_t0 = os.pread(out_fd, len(t0_bytes), doff + parsed["tensors"][0]["rel_offset"])
assert got_t0 == t0_bytes
# Check t1 sliced Q2_0 experts [1, 3, 4, 7] bit-identical
expected_t1 = b"".join(t1_exp_slices[i] for i in keep[0])
got_t1 = os.pread(out_fd, len(expected_t1), doff + parsed["tensors"][1]["rel_offset"])
assert got_t1 == expected_t1
# Check t2 sliced router rows [1, 3, 4, 7] bit-identical
expected_t2 = b"".join(t2_exp_slices[i] for i in keep[0])
got_t2 = os.pread(out_fd, len(expected_t2), doff + parsed["tensors"][2]["rel_offset"])
assert got_t2 == expected_t2
finally:
os.close(out_fd)
print("slice_qwen4exp_gguf.py self-test: PASS")
def main():
p = argparse.ArgumentParser(description="Slice experts from Qwen3.8-Flash-Next GGUF shards.")
p.add_argument("--self-test", action="store_true", help="Run built-in unit test and exit.")
p.add_argument("--in-gguf", help="Input GGUF shard path.")
p.add_argument("--out-gguf", help="Output GGUF shard path.")
p.add_argument("--keep-json", help="JSON file mapping layer index ('0'..'47') to list of kept expert IDs.")
p.add_argument("--old-experts", type=int, default=512)
args = p.parse_args()
if args.self_test:
run_self_test()
return
if not (args.in_gguf and args.out_gguf and args.keep_json):
p.error("--in-gguf, --out-gguf, and --keep-json are required unless --self-test is set")
raw_map = json.loads(Path(args.keep_json).read_text())
keep_map = {int(k): sorted(int(x) for x in v) for k, v in raw_map.items()}
res = slice_gguf_shard(args.in_gguf, args.out_gguf, keep_map, old_n_experts=args.old_experts)
print(json.dumps(res, indent=2))
if __name__ == "__main__":
main()