用 Python 重写 picoLM

def dequantize_q4_K(raw, n: int) -> np.ndarray:
    """144 bytes -> 256 weights (4.5 bpw)."""
    blk, nb = _blocks(raw, Q4_K, n)
    d, dmin = _fp16_field(blk, 0), _fp16_field(blk, 2)
    packed, qs = blk[:, 4:16], blk[:, 16:144]

    sc = np.empty((nb, 8), dtype=np.float32)
    mn = np.empty((nb, 8), dtype=np.float32)
    for j in range(8):
        if j < 4:
            sc[:, j] = packed[:, j] & 63
            mn[:, j] = packed[:, j + 4] & 63
        else:                       # six-bit fields straddle byte boundaries
            sc[:, j] = (packed[:, j + 4] & 0xF) | ((packed[:, j - 4] >> 6) << 4)
            mn[:, j] = (packed[:, j + 4] >> 4) | ((packed[:, j] >> 6) << 4)

    q = qs.reshape(nb, 4, 32)
    out = np.empty((nb, 4, 64), dtype=np.float32)
    d_, dm_ = d[:, :, None], dmin[:, :, None]
    out[:, :, :32] = d_ * sc[:, 0::2, None] * (q & 0xF) - dm_ * mn[:, 0::2, None]
    out[:, :, 32:] = d_ * sc[:, 1::2, None] * (q >> 4) - dm_ * mn[:, 1::2, None]
    return out.reshape(-1)

# picoLM / llama.cpp: rotate the pair (x[2i], x[2i+1])
for i in range(half):
    q0, q1 = qh[i * 2], qh[i * 2 + 1]
    qh[i * 2]     = q0 * cos[i] - q1 * sin[i]
    qh[i * 2 + 1] = q0 * sin[i] + q1 * cos[i]

# S03 / GPT-NeoX: rotate the pair (x[i], x[i + D/2])
out[:half] = x[:half] * cos - x[half:] * sin
out[half:] = x[half:] * cos + x[:half] * sin

float *model_forward(model_t *m, int token, int pos) {
    dequantize_row(embd_row, s->x, dim, w->type_token_embd);

    for (int l = 0; l < c->n_layers; l++) {
        rmsnorm(s->xb, s->x, s->attn_norm_w[l], dim);
        matmul(s->q, s->xb, lw->attn_q, dim, dim, lw->type_attn_q);
        /* … K, V … */
        rope(s->q, k_tmp, head_dim, n_heads, n_kv_heads, cos_pos, sin_pos);

        for (int d = 0; d < kv_dim; d++)
            key_pos_fp16[d] = fp32_to_fp16(k_tmp[d]);   /* the cache is fp16 */

        for (int h = 0; h < n_heads; h++) {
            /* online softmax — no att[] buffer is ever allocated */
            for (int t = 0; t <= pos; t++) { /* … */ }
        }
        /* … output projection, SwiGLU, residuals … */
    }
    rmsnorm(s->x, s->x, s->output_norm_w, dim);
    matmul(s->logits, s->x, w->output, dim, c->vocab_size, w->type_output);
    return s->logits;
}

def forward(self, token: int, pos: int) -> np.ndarray:
    c = self.cfg
    cos_pos, sin_pos = self.cos[pos], self.sin[pos]
    x = quant.dequantize_row(raw_row, c.n_embd, emb.qtype)

    for l in range(c.n_layers):
        xb = tensor.rmsnorm(x, self.attn_norm[l])
        q = tensor.matmul(xb, *self._w(f"blk.{l}.attn_q.weight"), c.n_embd, c.n_embd)
        k = tensor.matmul(xb, *self._w(f"blk.{l}.attn_k.weight"), c.n_embd, c.kv_dim)

        q = tensor.rope_interleaved(q, c.n_heads, c.head_dim, cos_pos, sin_pos)
        k = tensor.rope_interleaved(k, c.n_kv_heads, c.head_dim, cos_pos, sin_pos)

        self.key_cache[l, pos] = k.astype(np.float16)   # the cache is fp16
        self.val_cache[l, pos] = v.astype(np.float16)

        for h in range(c.n_heads):
            g = h // c.kv_mul                # the KV head this query shares
            scores = (kh[:, g, :] @ qh[h]) * inv
            out[h] = tensor.softmax(scores) @ vh[:, g, :]
        # … output projection, SwiGLU, residuals …

    x = tensor.rmsnorm(x, self.output_norm)
    return tensor.matmul(x, *self._w(self.output_name), c.n_embd, c.vocab_size)

for pos, t in enumerate(ids):          # prefill: one position at a time
    logits = model.forward(t, pos)

out_ids, pos = [], len(ids)
while len(out_ids) < max_tokens and pos < model.cfg.max_seq_len:
    row = grammar.mask_logits(logits, state, pieces) if grammar else logits
    nxt = sampler(row)
    if nxt == tok.eos_id:
        break
    out_ids.append(nxt)
    stream.write(dec.push(nxt))
    logits = model.forward(nxt, pos)
    pos += 1

# any Llama-architecture GGUF works — TinyLlama is the one picoLM targets
cd code
python -m picolm.run --model tinyllama-1.1b-chat.Q4_K_M.gguf \
                     --prompt "Once upon a time" -n 40

# grammar-constrained, so the output is JSON or nothing
python -m picolm.run --model tinyllama-1.1b-chat.Q4_K_M.gguf \
                     --prompt "Describe a cat as JSON" --json