FRSMASH v3.6

Fast Recursive State machine with Multi-slot Attention-recall System using Hybrid (SSM+GLA) architecture

SSM backbone(长程平滑) + GLA recall(精确召回) + SlowMemory(慢记忆) 三路融合架构 —— 训练快、推理显存恒定、长程外推稳定、具备 associative recall 能力。

目录

• 架构
• 架构优势
• 测试数据
• 环境安装
• 训练指南
• 推理指南
• 版本演进
• 文件说明

架构

input tokens
    │  + Positional Encoding
    ▼
┌─────────────────────────────────────────────┐
│  SSM backbone  (多槽 F-layer, fla HGRN)      │ ← 长程平滑, O(T·D), state 固定
│  12 层 × [MultiSlotFLayer + SwiGLU FFN]      │
└─────────────────────────────────────────────┘
    │ x_ash                    ┌──────────────────────────────────┐
    │                          │ SlowMemory (线性化, fla HGRN)     │ ← 慢记忆, O(T·D)
    │                          │ h_t = A(x)·h + B(x), y=α(x)·h+x  │
    │                          └──────────────────────────────────┘
    │ x_mem                              │
    │                          ┌──────────────────────────────────┐
    │                          │ Recall (fla multi-head GLA)       │ ← 精确召回, O(T·D·d_h)
    │                          │ S_t=exp(g)·S+k⊗v, o=q·S (content) │
    │                          └──────────────────────────────────┘
    │                                    │ x_recall
    ▼                                    ▼
  Gated Fusion(x_ash, x_mem) + x_emb + x_recall  →  Head → logits

三路分工:

• SSM backbone(多槽 F-layer):h = A·h + B,A∈(0,1],对数域 parallel scan(fla chunk_hgrn/fused_recurrent_hgrn)。长程平滑、信息有界、可并行。
• SlowMemory(线性化):h = A(x)·h + B(x) + 输出门。慢尺度记忆,backward 友好(fla autograd)。
• Recall(GLA):S_t = exp(g)·S_{t-1} + k⊗v,o = q·S。累积式 content addressing,解决纯 SSM 的 recall 弱点。multi-head(d_h=64),g 初始弱遗忘(bias=8,强保留)。

全 fla 实现:backbone/slowmem/recall 都用 flash-linear-attention 的高效 triton kernel,无 Python loop,训练/推理都快。

架构优势

核心优势:

1. recall 强:GLA 分支补上纯 SSM 的 associative recall 弱点(验证 induction PASS)。Mamba 做不到。
2. 生成显存恒定:三路 state 固定,生成任意长 context 显存不变(实测 276 MB,不随 context)。attention 的 KV cache 随序列线性涨。
3. 训练快:全 fla triton kernel + 自动 backward,无手写 backward(对比非线性 LSTM 的 backward 障碍)。
4. 长程外推稳:A∈(0,1] + exp(g)∈(0,1] 数学有界,128K(训练仅 512/768)外推不崩溃。

定位:v3.6 = Mamba-like SSM(效率) + RWKV-like GLA(recall) + SlowMemory(慢记忆) 融合,三路各取所长。

测试数据

训练(60M,d=432 8层,minimind 全量数据)

训练吞吐:125,000 tok/s(sl=512, bs=32, bf16, 单 RTX 4090)。

recall 能力(induction associative recall)

[序列 a1..aK marker a1..aK, 见 a_i 须输出后继]
v3.6: acc=1.000 PASS (step 800 学会)
对比: 纯 SSM(v3.4) acc=0.000 FAIL; cummax acc=0.058 FAIL; SKVR(slot EMA) acc=0.208 FAIL

只有累积式 content addressing(GLA)能过 induction,验证 recall 分支有效。

推理速度(60M)

生成速度不随 context 增长(state carry,每步 O(1))。

长程稳定性(外推)

context     NLL        PPL      maxTokNLL   drift(vs 2K)
  2K      10.558    38490       24.4        +0.00
  8K      10.225    27593       25.7        -0.33
 32K      10.134    25194       26.5        -0.42
128K      10.129    25064       27.2        -0.43
=> STABLE (训练 seq=512/768, 128K 是 170× 外推, PPL 不爆炸)

推理显存(关键优势)

context    生成显存(state carry)    attention 等价 KV cache
  100        276 MB                  1.3 MB
 1000        276 MB                 13.2 MB
 4000        276 MB                 52.7 MB
16000        276 MB                210.9 MB
            ↑ 恒定!                 ↑ 随 context 线性涨

生成显存 O(1)恒定(276 MB),attention 的 KV cache 随 context 线性增长。

prefill 峰值也恒定:generate/generate_cg 内部走 prefill_chunked(chunk=512),分块 forward 只保留 state、丢弃中间 activation,峰值显存 = 单块(实测 375 MB,不随 prompt 长度增长)。所以超长 prompt(16K+)也不会 OOM,prefill 与生成两阶段显存都恒定。

环境安装(任何机器)

1. 系统要求

• NVIDIA GPU(CUDA 11.8+ 或 12.x),建议 8GB+ 显存
• Python 3.10 ~ 3.12(不建议 3.13,fla 兼容性)

2. 安装依赖

# 创建环境
python -m venv .venv && .venv\Scripts\activate   # Windows
# 或: python3 -m venv .venv && source .venv/bin/activate  # Linux

# 安装
pip install torch                           # 按官网选对应 CUDA 版本
pip install flash-linear-attention          # fla (核心, triton kernel)
pip install jieba                            # OpenASHVoc 分词
pip install numpy tqdm                       # 数据/进度

fla 是关键依赖:提供 chunk_hgrn/fused_recurrent_hgrn(backbone/slowmem)和 chunk_gla(recall)的高效 triton kernel。纯 triton 实现,无需 nvcc 编译,Windows/Linux 通用。

若 fla 安装失败,可用 uv pip install flash-linear-attention(uv 处理依赖更稳)。

3. 准备数据与词表

需要 minimind 数据集 + OpenASH 词表:

data/
  pretrain_t2t_mini.jsonl    # ~1.27M 条 pretrain
  sft_t2t_mini.jsonl         # ~905K 条 sft
  open_ash_voc_agent.json    # 词表 (23005 tokens)

训练指南

配置数据路径

编辑 train/v36_train.py 顶部:

CACHE_DIR = '/path/to/data/cache'     # tokenized 缓存目录
OUT_DIR   = '/path/to/output'         # 权重输出目录
DEV       = 'cuda'

全量训练(pretrain + sft)

cd train
python v36_train.py --stage both --hidden 432 --layers 8 --heads 8 --pt_epochs 2 --sft_epochs 1

参数:

调试/冒烟

python v36_train.py --stage pretrain --max_steps 50   # 50 步验证

训练大规模(223M)

python v36_train.py --stage both --hidden 768 --layers 12 --pt_epochs 3 --sft_epochs 2
# ~48h (单 4090, 全量数据)

断点续训:自动从 out/v36_{stage}_latest.pth 恢复。

训练机制

• 优化器:AdamW(lr 5e-4 pretrain / 5e-5 sft, cosine schedule)
• 混合精度:bf16 autocast + GradScaler
• 梯度裁剪:1.0
• 缓存:首次 tokenize 后存 .pt,后续直接加载(省 jieba 分词)

推理指南

加载模型

import torch
from model.frsmash_v36 import FRSMASHv36
from infer.frsmash_v36_infer import generate, generate_cg

VS = 23005
model = FRSMASHv36(VS, 432, 8, 8, n_slots=4).cuda().eval()
model.load_state_dict(torch.load('out/v36_sft_final.pth', map_location='cuda', weights_only=False)['model'])

生成(支持 top-k / top-p / 重复抑制 / temperature)

prompt_ids = torch.tensor([[...]], device='cuda')   # tokenize 后的 ids

# eager 版 (调试友好, ~102 tok/s)
out = generate(model, prompt_ids, max_new=150,
               top_k=40, top_p=0.9, rep_penalty=1.3, temperature=0.9, eos=2)

# CUDA Graph 版 (生产用, ~625 tok/s, 显存恒定)
out = generate_cg(model, prompt_ids, max_new=150,
                  top_k=40, top_p=0.9, rep_penalty=1.3, temperature=0.9, eos=2)

三种生成方式对比

两者都:

• state carry(backbone + slowmem + recall 三路 state 传递),长 context 生成不重算
• 支持 top-k / top-p / repetition penalty / temperature 采样
• 生成显存恒定(276 MB,不随 context)
• chunked prefill(prefill_chunked,chunk=512):prefill 分块 forward + 三路 state carry(含 recall initial_state + pos_offset 位置编码偏移),峰值显存 = 单块,长 prompt 不 OOM

分词(OpenASHVoc)

from open_ash_voc import OpenASHVoc
voc = OpenASHVoc(agent_voc_path='data/open_ash_voc_agent.json')
ids = voc.encode("你好")          # text → ids
text = voc.decode(ids)            # ids → text

版本演进

v3.6 是完整解:v3.4 的速度(SSM 全 fla)+ v3.5 的 recall(GLA),fla 高效实现让两者兼得。

文件说明

FRSMASH-v3.6/
├── README.md                      # 本文档
├── model/
│   └── frsmash_v36.py             # 模型定义 (SSM + GLA recall + SlowMemory)
├── infer/
│   └── frsmash_v36_infer.py       # 推理模块 (generate / generate_cg + 采样)
└── train/
    └── v36_train.py               # 训练脚本 (pretrain → sft)

关键组件

• model/frsmash_v36.py:
  • MultiSlotFLayer:多槽 SSM F-layer(fla chunk_hgrn)
  • LinearSlowMemory:线性化慢记忆(fla chunk_hgrn)
  • GlaRecall:multi-head GLA 召回(fla chunk_gla + 手写 step for 生成)
  • FRSMASHv36:三路融合模型(forward 训练用,generate_step 推理用)
• infer/frsmash_v36_infer.py:
  • prefill_chunked:分块 prefill(chunk=512,三路 state carry,峰值显存恒定)
  • generate:eager 生成(fused_recurrent_hgrn 逐 token)
  • generate_cg:CUDA Graph 加速生成(6.5×,生产用)
  • _sample:top-k / top-p / repetition penalty / temperature
• train/v36_train.py:
  • CachedDS:tokenized 数据缓存(复用 minimind 缓存)
  • train:pretrain/sft 训练循环(bf16 + cosine lr + 断点续训)

引用

架构基于以下工作的思想融合:

• fla(flash-linear-attention):高效 triton kernel(chunk_hgrn, chunk_gla)
• Mamba(选择性 SSM):backbone 的 input-dependent 线性递推
• RWKV / GLA(gated linear attention):recall 分支的 content addressing
• Jamba(SSM + attention 混合):混合架构思路(v3.6 用更轻的 GLA 替代 full attention)

License

MIT

Citation

如果 FRSMASH v3.6 对你的研究或项目有帮助,欢迎引用:

@misc{frsmash_v36,
  title  = {FRSMASH v3.6: A Hybrid Architecture with SSM Backbone, GLA Recall and Linear SlowMemory},
  author = {dfytensor},
  year   = {2026},
  howpublished = {\url{https://www.modelscope.cn/dfytensor/FRSMASH-v3.6}},
  note   = {Three-branch fusion: multi-slot SSM F-layer (long-range, fla), 
            multi-head GLA recall (associative recall, fla), linear SlowMemory (slow memory). 
            Fully fla-based, O(1) inference memory, 128K extrapolation stable, induction recall PASS.}
}

关键特性一句话:SSM(长程效率) + GLA(精确召回) + SlowMemory(慢记忆) 三路融合,全 fla triton kernel,生成显存恒定 O(1),128K 长程外推稳定 —— 兼具 Mamba 的效率与 attention 的 recall 能力,且比 attention 推理显存更优。

欢迎在此基础上迭代(v3.7+),或应用到你的任务。如有问题/改进,欢迎 issue / PR。