daveycodez commited on
Commit
3c75909
·
verified ·
1 Parent(s): 2395323

Run the sliding-window layers with flash attention

Browse files
Files changed (2) hide show
  1. README.md +1 -0
  2. onnx/model_q4f16.onnx +2 -2
README.md CHANGED
@@ -28,6 +28,7 @@ This repository contains modified files. Compared with [onnx-community/LFM3-1.7B
28
  - **Smaller vocabulary:** the vocabulary was trimmed from 128,000 to 49,012 BPE tokens (plus the same 140 special tokens; 49,152 embedding rows). Every byte token and all special tokens (chat, tool-call and thinking tokens) are kept, and the merges stay consistent, so any text still tokenizes and decodes back exactly. Less common words split into more tokens, and non-English text needs more tokens. The embedding / LM head, `tokenizer.json`, `config.json` and `generation_config.json` (special token ids) changed to match.
29
  - **Shorter rotary caches:** the cos/sin caches were cut from 262,144 to 16,384 positions. Sequences longer than 16,384 tokens fail with a clear error instead of running.
30
  - **No second copy of the embedding at load:** in the original graph, a `Reshape` feeds the shared embedding / LM-head weights to the LM head, which makes onnxruntime-web copy that tensor while creating the session. The graph now gives the LM head the same weights without the `Reshape`. Outputs are identical.
 
31
  - **Calibration data:** about 36k tokens from the WikiText-2 training split plus chat replies generated by the fp16 model. No calibration text is included in the files.
32
 
33
  The chat template, `tokenizer_config.json` and `special_tokens_map.json` are unchanged.
 
28
  - **Smaller vocabulary:** the vocabulary was trimmed from 128,000 to 49,012 BPE tokens (plus the same 140 special tokens; 49,152 embedding rows). Every byte token and all special tokens (chat, tool-call and thinking tokens) are kept, and the merges stay consistent, so any text still tokenizes and decodes back exactly. Less common words split into more tokens, and non-English text needs more tokens. The embedding / LM head, `tokenizer.json`, `config.json` and `generation_config.json` (special token ids) changed to match.
29
  - **Shorter rotary caches:** the cos/sin caches were cut from 262,144 to 16,384 positions. Sequences longer than 16,384 tokens fail with a clear error instead of running.
30
  - **No second copy of the embedding at load:** in the original graph, a `Reshape` feeds the shared embedding / LM-head weights to the LM head, which makes onnxruntime-web copy that tensor while creating the session. The graph now gives the LM head the same weights without the `Reshape`. Outputs are identical.
31
+ - **Flash attention for the sliding-window layers:** the five sliding-window layers (window 256) get their window as an attention bias instead of `GroupQueryAttention`'s `local_window_size`. onnxruntime-web runs `local_window_size` layers without flash attention, which keeps a scores buffer of 16 heads × prompt length² (261 MB per layer for a 2,859-token prompt). The mask is the same, so outputs differ only by float rounding: on WebGPU, next-token loss over a 2,000-token prompt went from 3.0031 to 3.0009 and from 3.4849 to 3.4882 on two texts, and the difference is no larger inside the window than before it, where only the attention kernel changed. On a Mac, peak GPU buffers for a 2,859-token prompt went from 1,575 MB to 1,342 MB, and prompts that long process about 5% faster.
32
  - **Calibration data:** about 36k tokens from the WikiText-2 training split plus chat replies generated by the fp16 model. No calibration text is included in the files.
33
 
34
  The chat template, `tokenizer_config.json` and `special_tokens_map.json` are unchanged.
onnx/model_q4f16.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:6db32445ad347cd482ee6a59399ba5a24375f0db56669b481ca1c1f890330fd6
3
- size 296301
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a7d648f8cae27951cd2acc161af8f49cab5c50de1e97911e46bd73635f16c33f
3
+ size 298410