本文带你精读 Llama-3.2-1B-Instruct-singprobe 项目的核心源码文件 modeling_sing_probe.py，看懂 SingProbe 流式安全探针是如何实现的：它以不到 4.2M 的参数量和小于 0.5% 的解码开销，复用基座模型 Llama-3.2-1B-Instruct 的隐状态，在每个 token 上输出查询意图、响应不安全性和幻觉风险三类分数。文章将围绕三大核心模块展开——滑动窗口注意力（sliding window attention）、GQA 分组查询注意力与轻量分类器，帮你在 15 分钟内掌握其实现原理。

一、项目速览：一个只有 4.2M 参数的"安全探针"

SingProbe 是一个内在流式护栏（intrinsic streaming guardrail）：它不运行独立的安全模型，而是"搭便车"于基座模型生成过程中，对每个 token 实时打分。

关键配置来自 config.json：

📌 三层隐状态拼接后输入维度为2048 × 3 = 6144，这正是配置属性input_size的计算方式（见 configuration_sing_probe.py 中的input_size属性）。

二、源码地图：两种探针，一个基类

整个文件不到 150 行，结构非常清晰：

• SingProbePreTrainedModel（modeling_sing_probe.py#L13-L31）：基类，负责权重初始化（nn.Linear用std=0.02的正态分布，RMSNorm权重初始化为 1），以及最重要的输入校验——
• SingProbeMlpModel（modeling_sing_probe.py#L34-L57）：MLP 版探针，两层全连接 + 激活函数，结构最简单；
• SingProbeAttnModel（modeling_sing_probe.py#L60-L142）：注意力版探针，即本仓库config.json中auto_map指定的默认架构，也是本文主角。

输入校验逻辑值得留意（modeling_sing_probe.py#L25-L31）：如果输入特征维度不等于hidden_size × 取样层数，直接抛出ValueError，这保证了探针与基座模型的严格配对。

三、滑动窗口注意力：为什么、怎么实现

3.1 为什么要用滑动窗口

探针需要在流式生成中逐 token 评估，上下文可能很长。若每个 token 都对全部历史做因果注意力，代价随序列长度平方增长。SingProbe 采用sliding_window = 2048：每个 query 位置最多只回看最近 2048 个位置，在安全性与延迟之间取得平衡。

3.2 分块实现的三步拆解

核心方法是_sliding_window_attention（modeling_sing_probe.py#L80-L109），思路是把长序列切成window大小的块逐块计算：

1. 分块：for start in range(0, seq_len, window)，每次处理[start, end)区间的 query；
2. 构造掩码：对每个 query 位置，只允许关注满足0 <= q_pos - k_pos < window的 key 位置——既保证因果性（不看未来），又限制窗口（不看太远）：

relative_positions = query_positions[:, None] - key_positions attention_mask = (relative_positions >= 0) & (relative_positions < window)

3. 块内 Key 左移一格：key_start = max(0, start - window + 1)，让每个块的 key 范围比 query 块向前多覆盖一个位置，保证窗口边界处注意力不"漏看"相邻上下文；各块结果最后用torch.cat拼接。

💡 注意forward中的分支判断（modeling_sing_probe.py#L126-L135）：只有当sliding_window < seq_len时才走滑窗路径，短序列直接使用 PyTorch 原生的F.scaled_dot_product_attention(..., is_causal=True)一次算完，避免不必要的分块开销。这是很好的工程习惯：快路径优先。

四、GQA 分组查询注意力：8 个 Query 共享 1 组 K/V

观察__init__中的投影层定义（modeling_sing_probe.py#L72-L74）：

self.proj_q = nn.Linear(config.input_size, self.projection_size, bias=False) # 512 = 8 × 64 self.proj_k = nn.Linear(config.input_size, self.head_dim, bias=False) # 64 self.proj_v = nn.Linear(config.input_size, self.head_dim, bias=False) # 64

• Query 投影到8 × 64 = 512维，reshape 成8 个头；
• Key/Value 只投影到 64 维，即1 个头；
• 8 个 Query 头共享同一组 K/V——这就是GQA（Grouped Query Attention）的典型形态，介于 MHA（多头各自独立）与 MQA（全共享）之间。

实际计算时通过enable_gqa=self.num_attention_heads > 1交给 PyTorch 的scaled_dot_product_attention原生支持，一行代码完成分组复用。

为什么对探针特别合适？探针的输入是多层隐状态拼接（6144 维），K/V 共享可以大幅减少参数与 KV 存储，这正是它能把解码开销压到 0.5% 以下的关键设计之一。全模型约 4.2M 参数中，proj_q（约 3.1M）占了大头，其余部分非常轻量。

五、分类器：残差 + RMSNorm + 单层线性

注意力算完后，forward的最后四行（modeling_sing_probe.py#L136-L139）完成了从"上下文表示"到"10 维分数"的收尾：

context = context.transpose(1, 2).contiguous() context = context.view(batch_size, seq_len, self.projection_size) hidden_states = self.o_proj(context) + query_features # 残差连接 logits = self.classifier(self.norm(hidden_states)) # RMSNorm + Linear(512 → 10)

几个设计点：

• 残差连接：o_proj(context) + query_features，把原始查询特征加回去，类似 Transformer 中的残差思想，让梯度更稳定；
• RMSNorm：对 512 维投影空间做归一化（eps=1e-6），稳定分类前的数值尺度；
• 单层线性分类器：512 → 10维 logits，输出 8 类意图 + 不安全 + 幻觉，封装为标准的TokenClassifierOutput返回，方便上层推理框架按label_0–label_9取分。

六、动手验证：加载与运行 SingProbe

集成好的推理服务中，通过--probe-ckpt参数挂上探针即可（命令摘自 README.md）：

python -m sglang.launch_server \ --model-path meta-llama/Llama-3.2-1B-Instruct \ --probe-ckpt inclusionAI/Llama-3.2-1B-Instruct-singprobe \ --port 30000

生成的每个 token 会附带一个分数字典（label_0–label_9）。请务必使用精确配对的基座/探针组合：meta-llama/Llama-3.2-1B-Instruct+ 本仓库检查点。本地可获取仓库：

git clone https://gitcode.com/hf_mirrors/inclusionAI/Llama-3.2-1B-Instruct-singprobe

七、小结：一张表看懂实现要点

✅ 一句话总结：SingProbe 用多层隐状态拼接输入 + 滑动窗口 GQA + 极简分类头的组合，把一个流式安全评估任务做进了 4.2M 参数和 0.5% 解码开销里——这是"轻量探针"设计的一个教科书级范例。