File size: 6,134 Bytes
c8b4487
 
 
 
 
 
 
 
fac062c
 
 
c8b4487
 
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
c8b4487
fac062c
 
 
 
 
 
 
 
c8b4487
9f257c6
 
 
fac062c
c8b4487
fac062c
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
---
base_model: meta-llama/Llama-3.1-8B-Instruct
library_name: peft
pipeline_tag: text-generation
tags:
- base_model:adapter:meta-llama/Llama-3.1-8B-Instruct
- lora
- transformers
license: apache-2.0
language:
- en
---

# SSA-Llama-3.1-8B-Instruct

**Llama 3.1 8B Instruct with content-adaptive sparse attention for efficient long-context inference.**

SSA dynamically selects relevant regions of long context instead of requiring selected attention layers to attend densely across the entire sequence.

The goal is simple:

> **Preserve long-context capability while making attention increasingly efficient as context grows.**

This repository contains the primary **SSA-Llama-3.1-8B-Instruct checkpoint** trained in the SSA research.

## Highlights

* Built on **Meta Llama 3.1 8B Instruct**
* Content-adaptive sparse attention
* Long-context retrieval evaluated from **16K to 128K**
* Sparse-attention operator evaluated up to **1M tokens**
* Strong single-fact retrieval through **64K**
* Increasing computational advantage as context length grows
* Designed for research into practical long-context LLM inference

## Long-Context Retrieval

Using the current SSA sparse-attention backend:

| Context | Single-Fact NLL | Multi-Fact NLL |
| ------: | --------------: | -------------: |
|     16K |         0.00001 |        0.00006 |
|     32K |         0.00005 |        0.00114 |
|     48K |         0.00026 |        0.00237 |
|     64K |         0.00015 |        0.06979 |

**First-answer-token Top-1 remained 100% across all of these diagnostic cases.**

Single-fact retrieval remained extremely strong through 64K. Multi-fact retrieval was also strong through 48K, with increased difficulty appearing at 64K.

These numbers are controlled teacher-forced retrieval diagnostics. Autoregressive Needle-in-a-Haystack results will be added separately after evaluation is complete.

## Why Sparse Attention?

Standard dense causal attention becomes increasingly expensive as context grows.

SSA limits the amount of historical context processed by its sparse attention layers while still selecting information based on content.

Conceptually:

```text
Long Context
     ↓
Find Relevant Historical Regions
     ↓
Local Context + Selected Context
     ↓
Sparse Attention
     ↓
Generation
```

For the currently evaluated configuration, the nominal semantic attention support is approximately **8,196 keys per query**.

### Semantic Attention Arithmetic

The ideal semantic QK/PV arithmetic reduction grows rapidly with context length:

| Context | Dense causal avg. keys/query | SSA fanout cap | Attention arithmetic reduction | QK/PV work saved |
| ------: | ---------------------------: | -------------: | -----------------------------: | ---------------: |
|     32K |                       16,384 |          8,196 |                      **2.00×** |        **50.0%** |
|     64K |                       32,768 |          8,196 |                      **4.00×** |        **75.0%** |
|    128K |                       65,536 |          8,196 |                      **8.00×** |        **87.5%** |
|    256K |                      131,072 |          8,196 |                     **15.99×** |       **93.75%** |
|    512K |                      262,144 |          8,196 |                     **31.98×** |       **96.87%** |
|      1M |                     ~500,000 |          8,196 |                     **61.01×** |       **98.36%** |

These numbers represent **semantic attention arithmetic**, not total model FLOPs or end-to-end speedup.

## Measured Sparse-Attention Scaling

The custom SSA attention operator shows increasing benefit as context grows:

| Context | Relative operator speed |
| ------: | ----------------------: |
|     64K |                   ~1.0× |
|    128K |                   ~1.9× |
|    256K |                   ~3.7× |
|    512K |                   ~7.3× |
|      1M |              **~14.2×** |

The important distinction is that these are **attention-operator measurements**, not end-to-end LLM speedups.

The results show the central systems motivation behind SSA: sparse attention becomes increasingly valuable as context length grows.

## Model Details

* **Model:** `onlyaady/SSA-Llama-3.1-8B-Instruct`
* **Base model:** `meta-llama/Llama-3.1-8B-Instruct`
* **Developer:** Hrudayaditya Jallu
* **Affiliation:** Independent research
* **Checkpoint:** SSA-Llama-3.1-8B-Instruct
* **Format:** PEFT / LoRA
* **Language:** English
* **Status:** Research release

## Using the Model

This repository currently contains the trained PEFT checkpoint.

SSA also relies on a custom sparse-attention runtime. Loading this adapter through a standard Transformers pipeline alone will **not reproduce the sparse-attention behavior or performance measurements above**.

A reproducible inference/runtime package will be released separately.

This checkpoint is currently intended for research in:

* long-context LLM inference
* sparse attention
* content-based routing
* retrieval
* efficient inference
* quality/compute tradeoffs

## Current Limitations

SSA is in active research.

A few important limitations remain:

* only a subset of transformer layers currently use the SSA sparse path
* the complete model is therefore not yet fully subquadratic
* multi-fact retrieval becomes more difficult at larger context lengths, the model was only trained to support 16k.
* 1M-token speed results are operator-level measurements
* additional free-generation evaluations are still in progress

This checkpoint should currently be treated as a **research release rather than a production replacement for standard Llama 3.1 inference**.

## Research

This model is part of my research project.

Manuscript in preparation.

## Citation

```bibtex
@misc{jallu2026ssa,
  title  = {Sub-Quadratic Sparse Attention (SSA) for Long-Context LLMs},
  author = {Jallu, Hrudayaditya},
  year   = {2026},
  note   = {Manuscript in preparation}
}
```

## Contact for queries
Mail - [hjallu@umass.edu](mailto:hjallu@umass.edu)

## Framework

* PyTorch
* Hugging Face Transformers
* PEFT 0.19.1
* Meta Llama 3.1 8B Instruct