企业官网建设流程全解析

K2-Horizon-3.7B-4bit 是 IFM K2-Horizon 系列 3.7B 稠密模型在 Apple 芯片上的 4bit 量化 MLX 版本，磁盘占用仅 4.1 GB。这篇文章带你做代码级解析：核心文件 k2_horizon.py 中 GroupedRMSNorm 分组归一化、GQA 分组查询注意力、MoVA（Mixture-of-Values）值专家路由三大组件的实现原理，以及用 mlx-lm 快速加载这个模型的方法。

一、仓库速览：K2-Horizon-3.7B-4bit 参数与量化方案

这个仓库是 IFM/K2-Horizon-3.7B 的混合精度 4bit MLX 量化版（由 bf16 版转换而来）：

• 等效 6.5 bits/weight，权重文件 model.safetensors 共 4.1 GB；
• 短 prompt 峰值内存4.3 GB，16 GB 的 Mac 即可运行；
• 实测（Mac Studio M4 Max 128GB）生成速度109.6 tok/s，支持512K上下文。

config.json 关键参数一览：

一个有意思的细节：model.safetensors.index.json 的 metadata 显示参数总量约 5.06B（含 250624 词表的 embedding 与 lm_head），"3.7B" 是官方命名。

量化策略（config.json的quantization字段，affine 整数量化）：

• MLP 权重（gate_proj/up_proj/down_proj）：4-bit，group_size 32——参数最多，压得最狠；
• 其余全部（注意力 Q/K/V/O、embeddings、lm_head）：8-bit，group_size 64——更敏感，保留精度。

README.md 给了对比数据：纯 4bit（全部 group 64）会让 WikiText-2 困惑率上升 +21.5%，而这套混合精度方案只有+4.4%（18.324 vs bf16 的 17.552）。

注意：K2-Horizon 是一个模型家族，更大的成员使用 MoE + MoVA 稀疏层。k2_horizon.py 是全家族共用的 MLX 移植版（源自官方 PyTorch 实现modeling_k2_horizon.py，在 fp32 下逐层校验到 4e-6 相对误差）。而这份 3.7B checkpoint 是稠密模型：mlp_only_layers覆盖 0–35 全部层，MoVA 分支处于"就绪但未启用"状态。

二、GroupedRMSNorm 原理：2560 维为什么分 2 组归一化

普通 RMSNorm 对整条向量做一次归一化；GroupedRMSNorm把hidden_size切成layernorm_num_groups个等长分组，每组独立归一化，再乘一个全宽的缩放权重。本模型中 groups=2，即 2560 = 2 × 1280。

核心实现见 k2_horizon.py 的GroupedRMSNorm，分三步：

1. 先转fp32并 reshape 成(..., 2, 1280)——归一化只在 fp32 下做，避免低精度误差；
2. 调用融合内核mx.fast.rms_norm(y, None, eps)（此时不传权重）；
3. reshape 回原形状，乘全宽的self.weight，再转回输入精度。

源码里还有一条groups==1 的快速路径：直接mx.fast.rms_norm(x, weight, eps)，让内核一次完成归一化加乘权重。但注释特意说明：本 checkpoint 的所有归一化都走 groups=2 分支，不会碰快速路径。

它在代码里出现三处，全部使用layernorm_num_groups=2：

• DecoderLayer的input_layernorm与post_attention_layernorm（k2_horizon.py）；
• K2HorizonModel的最终输出归一化（k2_horizon.py）。

配合 Pre-norm 残差结构（k2_horizon.py），每层计算逻辑非常直白：

h = x + self_attn(input_layernorm(x), mask, cache) h = h + mlp(post_attention_layernorm(h))

三、GQA 实现：32 个 Q 头共享 8 组 KV 的代码路径

GQA（Grouped-Query Attention）介于标准 MHA 与 MQA 之间：Q 有 32 个头，K/V 只有 8 个头，每个 KV 头被 4 个 Query 头共享。好处是KV 缓存直接缩小到 1/4，长上下文推理更省内存。

Attention类（k2_horizon.py）的关键实现点：

1. 投影层：q_proj2560→4096（32×128）；k_proj/v_proj2560→1024（8×128）；o_proj4096→2560。attention_bias=false，全部无偏置；
2. QK 归一化：本 checkpointquery_key_norm=false，所以不创建q_norm/k_norm（代码中保留了"按头分组"的 GroupedRMSNorm 路径，供家族中需要它的配置使用）；
3. RoPE：traditional=False的 rotate-half 布局，基频 1e7 来自rope_parameters；带缓存时位置号直接取自cache.offset，无需额外传position_ids；
4. 注意力核心：调用 mlx-lm 的融合scaled_dot_product_attention（k2_horizon.py），原生支持 GQA，不需要把 KV 头repeat_kv复制成 32 份；
5. KV 缓存：使用 mlx-lm 的KVCache对象，cache.update_and_fetch(keys, values)一步完成写入和读取。

📌KV 缓存有多大？每层 K+V = 8 头 × 128 维 × 2（K 和 V）× 2 字节（bf16）= 4 KB/token，乘 36 层 ≈144 KB/token。128K 上下文约 18 GB，超长上下文建议配合--max-kv-size与 KV 量化--kv-bits 8控制内存。

四、MoVA 架构：值张量如何走专家路由

MoVA（Mixture-of-Values）是 K2-Horizon 的核心创新：普通注意力里 V 张量来自单个v_proj；在 MoVA 层里，V 由一组"值专家"路由混合产生——相当于把 MoE 的思想从 FFN 延伸到了注意力。

4.1 Router：五步路由算法

Router类（k2_horizon.py）被 MoE 块和 MoVA 共用，流程五步：

1. logits = x @ W.T（只用权重、不用偏置，按模型精度计算）；
2. 转fp32后算sigmoid（本家族配置）或softmax得到 scores；
3. 用scores + bias取 top-k——bias 只决定"选谁"，不影响权重，这点搞错输出会差很远；
4. 取被选专家的未加偏置分数作为权重；
5. 可选归一化（和为 1），再乘router_scaling_factor（本模型为 1.0）。

两个移植细节值得注意：torch.topk换成mx.argpartition（不排序，因为专家输出最后会求和，顺序无关）；权重全程保持 fp32（PyTorch 版会先转 bf16），结果只多不出错。

4.2 SwitchValueExperts：一次内核算完所有命中专家

SwitchValueExperts（k2_horizon.py）负责值的混合：

• 参考实现里 64 个无偏置值专家各是一个独立nn.Linear，PyTorch 原版用 Python 循环逐个专家处理、index_add_累加——这在 MLX 的惰性图里又慢又难支持；
• 移植后所有专家权重被压成一个[64, 1024, 2560]张量（由sanitize完成，见下一节），一次gather_mm内核调用跑完全部选中专家；
• 当 (token, expert) 对数 ≥ 64 时（即长 prompt 预填充阶段），先按专家排序（_gather_sort）让内核顺序读取权重，算完再散回原序；逐 token 解码时跳过排序；
• SiLU 在加权求和之前对每个专家输出单独施加（顺序与参考实现一致），加权和在 fp32 下完成再转回原精度。

4.3 在 Attention 中，MoVA 分支长什么样？

关键只有两行（k2_horizon.py）：

if self.mova: values = self.v_experts(x, *self.v_router(x)) # 路由出的值专家混合 else: values = self.v_proj(x) # 普通投影

MoVA 层里没有v_proj，取而代之的是v_router+v_experts。由于值张量仍是 8×128 的形状、按 token 从层输入算出，它进入 KV 缓存的方式与普通 V 完全一致——MoVA 不需要任何特殊缓存，这是设计上的一个优雅之处。

层的判定规则在DecoderLayer（k2_horizon.py）：层不在mlp_only_layers中、且num_experts > 0、且层号命中decoder_sparse_step节奏时，才是"稀疏层"（MoE FFN + MoVA 注意力）。本 3.7B checkpoint 的mlp_only_layers覆盖全部 36 层，所以mova恒为False，走普通路径。

五、MoE 块与 sanitize：稀疏层的完整链路

SparseMoeBlock（k2_horizon.py）是 FFN 侧的稀疏部分：

• gate：复用上文Router（本模型norm_topk_prob=true时会归一化权重）；
• switch_mlp：SwitchGLU路由专家（SwiGLU 激活）；
• 共享专家：num_shared_experts > 0时是一个更宽的 MLP，输出直接相加、不加门控。

前向流程：gate出(indices, scores)→ 专家计算 → fp32 加权和 → 加共享专家输出。本 checkpointnum_experts=0、moe_intermediate_size=0，此类根本不会被实例化。

权重加载由Model.sanitize（k2_horizon.py）完成三件事：

1. 把 MoE 层中每个专家的独立权重堆叠成一个[num_experts, out, in]张量——gather_mm单张量推理的前提；
2. 同样堆叠 MoVA 的v_experts权重；
3. 丢弃不需要加载的rotary_emb.inv_freq（RoPE 频率运行时现算），词表共享时丢弃lm_head。

该函数幂等，重复执行也安全。

六、mlx-lm 运行指南：加载 4bit 模型的一键步骤

由于 mlx-lm 尚未内置k2_horizon架构，本仓库直接携带 MLX 模型代码 k2_horizon.py，并由 config.json 的model_file字段指向它，官方 mlx-lm 开箱即用：

pip install -U mlx-lm git clone https://gitcode.com/hf_mirrors/mlx-community/K2-Horizon-3.7B-4bit mlx_lm.generate --model K2-Horizon-3.7B-4bit --trust-remote-code \ --prompt "用两句话解释 mixture-of-experts。" --max-tokens 2048

📌 使用提示：

• 新版 mlx-lm 必须加--trust-remote-code，它会加载并运行仓库内的自定义模型文件（文件注释详尽，建议先读）；
• 模型会先思考后作答：输出包在<ifm|think> … </ifm|think>标签内，max_tokens要预留空间；思考力度可在 chat_template.jinja 的reasoning_effort参数设为high/medium/low；
• 分词器与生成配置见 tokenizer.json、tokenizer_config.json、generation_config.json；
• mlx_lm.chat与mlx_lm.server同样支持，但 server 目前会把思考文本直接放在content字段里。

七、常见问题 FAQ

Q1：这个 3.7B 模型启用了 MoVA 吗？没有。config.json中mova_num_experts=0、num_experts=0，36 层全部是稠密 MLP。MoVA/MoE 代码是为家族中稀疏成员准备的，本 checkpoint 只走稠密路径。

Q2：为什么 4bit 只掉 4.4% 困惑率？靠混合精度方案：占参数大头的 MLP 用 4bit（group 32），敏感的注意力/embedding/lm_head 保留 8bit（group 64）。若全部 4bit，困惑率会恶化 21.5%。

Q3：为什么 Router 的权重不参与量化？Router是自定义模块而非nn.Linear，mlx-lm 的量化规则不会触碰它，路由权重保持全精度——保证"选哪个专家"这个决策不受量化误差干扰。

Q4：512K 上下文真的能跑吗？max_position_embeddings=524288支持它，但 KV 缓存约 144 KB/token（bf16），128K 上下文就约 18 GB。实际使用中用--max-kv-size限制缓存、--kv-bits 8压缩 KV，按需取舍。

总结：k2_horizon.py 仅 628 行，却完整覆盖了 K2-Horizon 家族的三大架构特征——用 fp32 分组的 GroupedRMSNorm 保证归一化精度、原生 GQA + 融合注意力内核压缩 KV 缓存、gather_mm 单次内核实现的 MoVA 值专家路由。对这个 3.7B 4bit checkpoint 而言，实际生效的是稠密路径 + 混合精度量化，也正是这三点组合让它以 4.1 GB 的体积在 16 GB Mac 上跑出 109.6 tok/s 的速度。