A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 20:00 UTC
[Model] Add block-local attention and YaRN for local layers to Gemma3 (#39823)
Signed-off-by: Philip Monk <169196560+philip-essential@users.noreply.github.com>
philip-essential committed
Apr 21, 2026 at 23:34 UTC
123674879e8ebb714b9a88c853806826b95b8f81
7 files changed
+540
-8
docs/models/supported_models.md
+1
@@ -472,6 +472,7 @@ th {
472
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
473
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
474
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
475
+| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
476
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
477
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
478
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+4
@@ -518,6 +518,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
518
extras={"tiny-random": "tiny-random/qwen3-next-moe"},
519
min_transformers_version="4.56.3",
520
),
521
+ "Rnj1ForCausalLM": _HfExamplesInfo(
522
+ "EssentialAI/rnj-1-instruct",
523
+ is_available_online=False,
524
+ ),
525
"RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),
526
"SarvamMoEForCausalLM": _HfExamplesInfo(
527
"sarvamai/sarvam-30b",
vllm/model_executor/layers/attention/attention.py
+6
@@ -336,6 +336,12 @@ class Attention(nn.Module, AttentionLayerBase):
336
)
337
cache_config.enable_prefix_caching = False
338
339
+ if extra_impl_args.get("chunk_lookback", -1) > -1:
340
+ assert self.attn_backend.get_name() == "TRITON_ATTN", (
341
+ f"Chunked attention with lookback requires the Triton backend, "
342
+ f"but got {self.attn_backend.get_name()}."
343
+ )
344
+
345
impl_cls = self.attn_backend.get_impl_cls()
346
self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an AttentionImpl subclass
347
num_heads,
vllm/model_executor/models/registry.py
+1
@@ -110,6 +110,7 @@ _TEXT_GENERATION_MODELS = {
110
"GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
111
"Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
112
"Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
113
+ "Rnj1ForCausalLM": ("rnj1", "Rnj1ForCausalLM"),
114
"Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
115
"Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"),
116
"Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
vllm/model_executor/models/rnj1.py
new
+470
@@ -0,0 +1,470 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# RNJ-1 model: Gemma3-based architecture with chunked (block-local) attention.
5
+# Chunked attention restricts local layers to attend within aligned blocks,
6
+# with lookback to one previous block.
7
+from collections.abc import Iterable
8
+from itertools import islice
9
+
10
+import torch
11
+from torch import nn
12
+from transformers import Gemma3TextConfig
13
+
14
+from vllm.compilation.decorators import support_torch_compile
15
+from vllm.config import CacheConfig, VllmConfig
16
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
17
+from vllm.logger import init_logger
18
+from vllm.model_executor.layers.activation import GeluAndMul
19
+from vllm.model_executor.layers.attention import Attention
20
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm
21
+from vllm.model_executor.layers.linear import (
22
+ MergedColumnParallelLinear,
23
+ QKVParallelLinear,
24
+ RowParallelLinear,
25
+)
26
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
+from vllm.model_executor.layers.quantization import QuantizationConfig
28
+from vllm.model_executor.layers.rotary_embedding import get_rope
29
+from vllm.model_executor.layers.vocab_parallel_embedding import (
30
+ ParallelLMHead,
31
+ VocabParallelEmbedding,
32
+)
33
+from vllm.model_executor.model_loader.weight_utils import (
34
+ default_weight_loader,
35
+ maybe_remap_kv_scale_name,
36
+)
37
+from vllm.sequence import IntermediateTensors
38
+from vllm.v1.attention.backend import AttentionType
39
+
40
+from .interfaces import SupportsLoRA, SupportsPP
41
+from .utils import (
42
+ AutoWeightsLoader,
43
+ extract_layer_index,
44
+ is_pp_missing_parameter,
45
+ make_empty_intermediate_tensors_factory,
46
+ make_layers,
47
+ maybe_prefix,
48
+)
49
+
50
+logger = init_logger(__name__)
51
+
52
+
53
+class Rnj1MLP(nn.Module):
54
+ def __init__(
55
+ self,
56
+ hidden_size: int,
57
+ intermediate_size: int,
58
+ hidden_activation: str,
59
+ quant_config: QuantizationConfig | None = None,
60
+ prefix: str = "",
61
+ ) -> None:
62
+ super().__init__()
63
+ self.gate_up_proj = MergedColumnParallelLinear(
64
+ hidden_size,
65
+ [intermediate_size] * 2,
66
+ bias=False,
67
+ quant_config=quant_config,
68
+ prefix=f"{prefix}.gate_up_proj",
69
+ )
70
+ self.down_proj = RowParallelLinear(
71
+ intermediate_size,
72
+ hidden_size,
73
+ bias=False,
74
+ quant_config=quant_config,
75
+ prefix=f"{prefix}.down_proj",
76
+ )
77
+ if hidden_activation != "gelu_pytorch_tanh":
78
+ raise ValueError(
79
+ "RNJ-1 uses `gelu_pytorch_tanh` as the hidden activation "
80
+ "function. Please set `hidden_act` and `hidden_activation` to "
81
+ "`gelu_pytorch_tanh`."
82
+ )
83
+ self.act_fn = GeluAndMul(approximate="tanh")
84
+
85
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
86
+ gate_up, _ = self.gate_up_proj(x)
87
+ x = self.act_fn(gate_up)
88
+ x, _ = self.down_proj(x)
89
+ return x
90
+
91
+
92
+class Rnj1Attention(nn.Module):
93
+ def __init__(
94
+ self,
95
+ config: Gemma3TextConfig,
96
+ hidden_size: int,
97
+ num_heads: int,
98
+ num_kv_heads: int,
99
+ head_dim: int,
100
+ max_position_embeddings: int,
101
+ cache_config: CacheConfig | None = None,
102
+ quant_config: QuantizationConfig | None = None,
103
+ attn_logits_soft_cap: float | None = None,
104
+ prefix: str = "",
105
+ ) -> None:
106
+ super().__init__()
107
+ self.config = config
108
+ self.hidden_size = hidden_size
109
+ tp_size = get_tensor_model_parallel_world_size()
110
+ self.total_num_heads = num_heads
111
+ assert self.total_num_heads % tp_size == 0
112
+ self.num_heads = self.total_num_heads // tp_size
113
+ self.total_num_kv_heads = num_kv_heads
114
+ if self.total_num_kv_heads >= tp_size:
115
+ assert self.total_num_kv_heads % tp_size == 0
116
+ else:
117
+ assert tp_size % self.total_num_kv_heads == 0
118
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
119
+ self.head_dim = head_dim
120
+ self.q_size = self.num_heads * self.head_dim
121
+ self.kv_size = self.num_kv_heads * self.head_dim
122
+ self.scaling = config.query_pre_attn_scalar**-0.5
123
+
124
+ self.qkv_proj = QKVParallelLinear(
125
+ hidden_size,
126
+ self.head_dim,
127
+ self.total_num_heads,
128
+ self.total_num_kv_heads,
129
+ bias=config.attention_bias,
130
+ quant_config=quant_config,
131
+ prefix=f"{prefix}.qkv_proj",
132
+ )
133
+ self.o_proj = RowParallelLinear(
134
+ self.total_num_heads * self.head_dim,
135
+ hidden_size,
136
+ bias=config.attention_bias,
137
+ quant_config=quant_config,
138
+ prefix=f"{prefix}.o_proj",
139
+ )
140
+
141
+ self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
142
+ self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
143
+
144
+ layer_idx = extract_layer_index(prefix)
145
+ layer_type = config.layer_types[layer_idx]
146
+ self.is_chunked = layer_type == "chunked_attention"
147
+ self.chunk_lookback = 1 if self.is_chunked else -1
148
+ sliding_window = config.sliding_window if self.is_chunked else None
149
+
150
+ # Initialize the rotary embedding.
151
+ # Expects v5-style rope_parameters keyed by layer type.
152
+ if layer_type in config.rope_parameters:
153
+ rope_parameters = config.rope_parameters[layer_type]
154
+ else:
155
+ rope_parameters = config.rope_parameters
156
+
157
+ self.rotary_emb = get_rope(
158
+ self.head_dim,
159
+ max_position=max_position_embeddings,
160
+ rope_parameters=rope_parameters,
161
+ is_neox_style=True,
162
+ )
163
+
164
+ self.attn = Attention(
165
+ self.num_heads,
166
+ self.head_dim,
167
+ self.scaling,
168
+ num_kv_heads=self.num_kv_heads,
169
+ cache_config=cache_config,
170
+ quant_config=quant_config,
171
+ attn_type=AttentionType.DECODER,
172
+ logits_soft_cap=attn_logits_soft_cap,
173
+ per_layer_sliding_window=sliding_window,
174
+ chunk_lookback=self.chunk_lookback,
175
+ prefix=f"{prefix}.attn",
176
+ )
177
+
178
+ def forward(
179
+ self,
180
+ positions: torch.Tensor,
181
+ hidden_states: torch.Tensor,
182
+ **kwargs,
183
+ ) -> torch.Tensor:
184
+ qkv, _ = self.qkv_proj(hidden_states)
185
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
186
+
187
+ q = q.unflatten(-1, (self.num_heads, self.head_dim))
188
+ q = self.q_norm(q)
189
+ q = q.flatten(-2, -1)
190
+ k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
191
+ k = self.k_norm(k)
192
+ k = k.flatten(-2, -1)
193
+
194
+ q, k = self.rotary_emb(positions, q, k)
195
+ attn_output = self.attn(q, k, v)
196
+ output, _ = self.o_proj(attn_output)
197
+ return output
198
+
199
+
200
+class Rnj1DecoderLayer(nn.Module):
201
+ def __init__(
202
+ self,
203
+ config: Gemma3TextConfig,
204
+ cache_config: CacheConfig | None = None,
205
+ quant_config: QuantizationConfig | None = None,
206
+ prefix: str = "",
207
+ ) -> None:
208
+ super().__init__()
209
+ self.hidden_size = config.hidden_size
210
+ self.self_attn = Rnj1Attention(
211
+ config=config,
212
+ hidden_size=self.hidden_size,
213
+ num_heads=config.num_attention_heads,
214
+ num_kv_heads=config.num_key_value_heads,
215
+ head_dim=config.head_dim,
216
+ max_position_embeddings=config.max_position_embeddings,
217
+ cache_config=cache_config,
218
+ quant_config=quant_config,
219
+ attn_logits_soft_cap=None,
220
+ prefix=f"{prefix}.self_attn",
221
+ )
222
+ self.hidden_size = config.hidden_size
223
+ self.mlp = Rnj1MLP(
224
+ hidden_size=self.hidden_size,
225
+ intermediate_size=config.intermediate_size,
226
+ hidden_activation=config.hidden_activation,
227
+ quant_config=quant_config,
228
+ prefix=f"{prefix}.mlp",
229
+ )
230
+ self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
231
+ self.post_attention_layernorm = GemmaRMSNorm(
232
+ config.hidden_size, eps=config.rms_norm_eps
233
+ )
234
+ self.pre_feedforward_layernorm = GemmaRMSNorm(
235
+ config.hidden_size, eps=config.rms_norm_eps
236
+ )
237
+ self.post_feedforward_layernorm = GemmaRMSNorm(
238
+ config.hidden_size, eps=config.rms_norm_eps
239
+ )
240
+
241
+ def forward(
242
+ self,
243
+ positions: torch.Tensor,
244
+ hidden_states: torch.Tensor,
245
+ residual: torch.Tensor | None,
246
+ **kwargs,
247
+ ) -> tuple[torch.Tensor, torch.Tensor]:
248
+ if residual is None:
249
+ residual = hidden_states
250
+ hidden_states = self.input_layernorm(hidden_states)
251
+ else:
252
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
253
+ hidden_states = self.self_attn(
254
+ positions=positions,
255
+ hidden_states=hidden_states,
256
+ **kwargs,
257
+ )
258
+ hidden_states = self.post_attention_layernorm(hidden_states)
259
+
260
+ hidden_states, residual = self.pre_feedforward_layernorm(
261
+ hidden_states, residual
262
+ )
263
+ hidden_states = self.mlp(hidden_states)
264
+ hidden_states = self.post_feedforward_layernorm(hidden_states)
265
+ return hidden_states, residual
266
+
267
+
268
+@support_torch_compile
269
+class Rnj1Model(nn.Module):
270
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
271
+ super().__init__()
272
+ config = vllm_config.model_config.hf_config
273
+ cache_config = vllm_config.cache_config
274
+ quant_config = vllm_config.quant_config
275
+ self.config = config
276
+ self.quant_config = quant_config
277
+
278
+ self.embed_tokens = VocabParallelEmbedding(
279
+ config.vocab_size,
280
+ config.hidden_size,
281
+ quant_config=quant_config,
282
+ prefix=f"{prefix}.embed_tokens",
283
+ )
284
+ self.start_layer, self.end_layer, self.layers = make_layers(
285
+ config.num_hidden_layers,
286
+ lambda prefix: Rnj1DecoderLayer(
287
+ config, cache_config, quant_config, prefix=prefix
288
+ ),
289
+ prefix=f"{prefix}.layers",
290
+ )
291
+ self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
292
+
293
+ normalizer = self.config.hidden_size**0.5
294
+ self.register_buffer("normalizer", torch.tensor(normalizer), persistent=False)
295
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
296
+ ["hidden_states", "residual"], config.hidden_size
297
+ )
298
+
299
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
300
+ return self.embed_tokens(input_ids) * self.normalizer
301
+
302
+ def forward(
303
+ self,
304
+ input_ids: torch.Tensor | None,
305
+ positions: torch.Tensor,
306
+ intermediate_tensors: IntermediateTensors | None,
307
+ inputs_embeds: torch.Tensor | None = None,
308
+ **kwargs,
309
+ ) -> torch.Tensor | IntermediateTensors:
310
+ if get_pp_group().is_first_rank:
311
+ if inputs_embeds is not None:
312
+ hidden_states = inputs_embeds
313
+ else:
314
+ hidden_states = self.embed_input_ids(input_ids)
315
+ residual = None
316
+ else:
317
+ assert intermediate_tensors is not None
318
+ hidden_states = intermediate_tensors["hidden_states"]
319
+ residual = intermediate_tensors["residual"]
320
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
321
+ hidden_states, residual = layer(
322
+ positions,
323
+ hidden_states,
324
+ residual,
325
+ **kwargs,
326
+ )
327
+ if not get_pp_group().is_last_rank:
328
+ return IntermediateTensors(
329
+ {"hidden_states": hidden_states, "residual": residual}
330
+ )
331
+ hidden_states, _ = self.norm(hidden_states, residual)
332
+ return hidden_states
333
+
334
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
335
+ stacked_params_mapping = [
336
+ # (param_name, shard_name, shard_id)
337
+ ("qkv_proj", "q_proj", "q"),
338
+ ("qkv_proj", "k_proj", "k"),
339
+ ("qkv_proj", "v_proj", "v"),
340
+ ("gate_up_proj", "gate_proj", 0),
341
+ ("gate_up_proj", "up_proj", 1),
342
+ ]
343
+ params_dict = dict(self.named_parameters())
344
+ loaded_params: set[str] = set()
345
+ for name, loaded_weight in weights:
346
+ if (
347
+ self.quant_config
348
+ and self.quant_config.get_name() == "gguf"
349
+ and name.endswith("norm.weight")
350
+ ):
351
+ loaded_weight -= 1
352
+
353
+ if self.quant_config is not None and (
354
+ scale_name := self.quant_config.get_cache_scale(name)
355
+ ):
356
+ param = params_dict[scale_name]
357
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
358
+ loaded_weight = loaded_weight[0]
359
+ weight_loader(param, loaded_weight)
360
+ loaded_params.add(scale_name)
361
+ continue
362
+
363
+ if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
364
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
365
+ if remapped_name is not None and remapped_name in params_dict:
366
+ param = params_dict[remapped_name]
367
+ weight_loader = getattr(
368
+ param, "weight_loader", default_weight_loader
369
+ )
370
+ weight_loader(param, loaded_weight)
371
+ loaded_params.add(remapped_name)
372
+ continue
373
+
374
+ for param_name, shard_name, shard_id in stacked_params_mapping:
375
+ if shard_name not in name:
376
+ continue
377
+ name = name.replace(shard_name, param_name)
378
+ if name.endswith(".bias") and name not in params_dict:
379
+ continue
380
+ if is_pp_missing_parameter(name, self):
381
+ continue
382
+ param = params_dict[name]
383
+ weight_loader = param.weight_loader
384
+ weight_loader(param, loaded_weight, shard_id)
385
+ break
386
+ else:
387
+ if name.endswith(".bias") and name not in params_dict:
388
+ continue
389
+ name = maybe_remap_kv_scale_name(name, params_dict)
390
+ if name is None:
391
+ continue
392
+ if is_pp_missing_parameter(name, self):
393
+ continue
394
+ param = params_dict[name]
395
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
396
+ weight_loader(param, loaded_weight)
397
+ loaded_params.add(name)
398
+
399
+ return loaded_params
400
+
401
+
402
+class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
403
+ packed_modules_mapping = {
404
+ "qkv_proj": [
405
+ "q_proj",
406
+ "k_proj",
407
+ "v_proj",
408
+ ],
409
+ "gate_up_proj": [
410
+ "gate_proj",
411
+ "up_proj",
412
+ ],
413
+ }
414
+
415
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
416
+ config = vllm_config.model_config.hf_config
417
+ quant_config = vllm_config.quant_config
418
+
419
+ super().__init__()
420
+ self.config = config
421
+ self.quant_config = quant_config
422
+ self.model = Rnj1Model(
423
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
424
+ )
425
+
426
+ self.lm_head = ParallelLMHead(
427
+ config.vocab_size,
428
+ config.hidden_size,
429
+ quant_config=quant_config,
430
+ prefix=maybe_prefix(prefix, "lm_head"),
431
+ )
432
+ if config.tie_word_embeddings:
433
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
434
+
435
+ self.logits_processor = LogitsProcessor(
436
+ config.vocab_size, soft_cap=config.final_logit_softcapping
437
+ )
438
+ self.make_empty_intermediate_tensors = (
439
+ self.model.make_empty_intermediate_tensors
440
+ )
441
+
442
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
443
+ return self.model.embed_input_ids(input_ids)
444
+
445
+ def forward(
446
+ self,
447
+ input_ids: torch.Tensor | None,
448
+ positions: torch.Tensor,
449
+ intermediate_tensors: IntermediateTensors | None = None,
450
+ inputs_embeds: torch.Tensor | None = None,
451
+ **kwargs,
452
+ ) -> torch.Tensor | IntermediateTensors:
453
+ hidden_states = self.model(
454
+ input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
455
+ )
456
+ return hidden_states
457
+
458
+ def compute_logits(
459
+ self,
460
+ hidden_states: torch.Tensor,
461
+ ) -> torch.Tensor | None:
462
+ logits = self.logits_processor(self.lm_head, hidden_states)
463
+ return logits
464
+
465
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
466
+ loader = AutoWeightsLoader(
467
+ self,
468
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
469
+ )
470
+ return loader.load_weights(weights)
vllm/v1/attention/backends/triton_attn.py
+3
@@ -458,6 +458,7 @@ class TritonAttentionImpl(AttentionImpl):
458
kv_sharing_target_layer_name: int | None = None,
459
sinks: torch.Tensor | None = None,
460
use_alibi_sqrt: bool = False,
461
+ chunk_lookback: int = -1,
462
) -> None:
463
self.num_heads = num_heads
464
self.head_size = head_size
@@ -492,6 +493,7 @@ class TritonAttentionImpl(AttentionImpl):
493
f"num_heads: {num_heads}."
494
)
495
self.use_alibi_sqrt = use_alibi_sqrt
496
+ self.chunk_lookback = chunk_lookback
497
self.supports_quant_query_input = current_platform.is_cuda()
498
499
self._kv_quant_mode = get_kv_quant_mode(kv_cache_dtype)
@@ -631,6 +633,7 @@ class TritonAttentionImpl(AttentionImpl):
633
kv_quant_mode=self._kv_quant_mode,
634
k_scale_cache=k_scale_cache,
635
v_scale_cache=v_scale_cache,
636
+ chunk_lookback=self.chunk_lookback,
637
)
638
639
return output
vllm/v1/attention/ops/triton_unified_attention.py
+55
-8
@@ -177,6 +177,8 @@ def kernel_unified_attention_2d(
177
stride_vs_blk=0,
178
stride_vs_slot=0,
179
stride_vs_head=0,
180
+ CHUNK_LOOKBACK: tl.constexpr = -1,
181
+ CHUNK_SIZE: tl.constexpr = -1,
182
):
183
q_block_global_idx = tl.program_id(0)
184
kv_head_idx = tl.program_id(1)
@@ -293,7 +295,11 @@ def kernel_unified_attention_2d(
295
# where q_abs = context_len + q
296
# The union of allowed key positions for this Q-block is:
297
# [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
296
- first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1
298
+ q_abs = context_len + qpos_lo
299
+ if CHUNK_LOOKBACK > -1:
300
+ first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE
301
+ else:
302
+ first_allowed_key = q_abs - SLIDING_WINDOW + 1
303
last_allowed_key = context_len + qpos_hi
304
# Convert to tile indices and clamp
305
tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
@@ -370,9 +376,19 @@ def kernel_unified_attention_2d(
376
query_abs_pos = context_len + query_pos[:, None]
377
seq_mask = seq_offset[None, :] <= query_abs_pos
378
373
- # Apply sliding window to base mask BEFORE mm_prefix OR.
374
- # Order must match FlexAttention: (causal AND sliding_window) OR mm_prefix
375
- if SLIDING_WINDOW > 0:
379
+ # Apply sliding window / chunked attention to base mask
380
+ # BEFORE mm_prefix OR.
381
+ # Order must match FlexAttention:
382
+ # (causal AND sliding_window) OR mm_prefix
383
+ if CHUNK_LOOKBACK > -1:
384
+ seq_mask = seq_mask & (
385
+ (
386
+ (context_len + query_pos[:, None]) // CHUNK_SIZE
387
+ - (seq_offset[None, :] // CHUNK_SIZE)
388
+ )
389
+ <= CHUNK_LOOKBACK
390
+ )
391
+ elif SLIDING_WINDOW > 0:
392
seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
393
394
# PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
@@ -559,6 +575,8 @@ def kernel_unified_attention_3d(
575
stride_vs_blk=0,
576
stride_vs_slot=0,
577
stride_vs_head=0,
578
+ CHUNK_LOOKBACK: tl.constexpr = -1,
579
+ CHUNK_SIZE: tl.constexpr = -1,
580
):
581
q_block_global_idx = tl.program_id(0)
582
kv_head_idx = tl.program_id(1)
@@ -681,7 +699,11 @@ def kernel_unified_attention_3d(
699
# where q_abs = context_len + q
700
# The union of allowed key positions for this Q-block is:
701
# [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
684
- first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1
702
+ q_abs = context_len + qpos_lo
703
+ if CHUNK_LOOKBACK > -1:
704
+ first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE
705
+ else:
706
+ first_allowed_key = q_abs - SLIDING_WINDOW + 1
707
last_allowed_key = context_len + qpos_hi
708
# Convert to tile indices and clamp
709
tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
@@ -761,9 +783,19 @@ def kernel_unified_attention_3d(
783
query_abs_pos = context_len + query_pos[:, None]
784
seq_mask = seq_offset[None, :] <= query_abs_pos
785
764
- # Apply sliding window to base mask BEFORE mm_prefix OR.
765
- # Order must match FlexAttention: (causal AND sliding_window) OR mm_prefix
766
- if SLIDING_WINDOW > 0:
786
+ # Apply sliding window / chunked attention to base mask
787
+ # BEFORE mm_prefix OR.
788
+ # Order must match FlexAttention:
789
+ # (causal AND sliding_window) OR mm_prefix
790
+ if CHUNK_LOOKBACK > -1:
791
+ seq_mask = seq_mask & (
792
+ (
793
+ (context_len + query_pos[:, None]) // CHUNK_SIZE
794
+ - (seq_offset[None, :] // CHUNK_SIZE)
795
+ )
796
+ <= CHUNK_LOOKBACK
797
+ )
798
+ elif SLIDING_WINDOW > 0:
799
seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
800
801
# PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
@@ -1046,6 +1078,8 @@ def unified_attention(
1078
kv_quant_mode: KVQuantMode = KVQuantMode.NONE,
1079
k_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
1080
v_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
1081
+ # Chunked attention: restrict attention to aligned blocks with lookback.
1082
+ chunk_lookback=-1,
1083
):
1084
assert causal, "Only causal attention is supported"
1085
assert q_descale is None, "Q scales not supported"
@@ -1093,6 +1127,15 @@ def unified_attention(
1127
# Tile sizes for prefill and decode. Gemma3 models use optimized values.
1128
# Note: tile size must be at least 32 for fp8 (element_size == 1).
1129
sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
1130
+
1131
+ # Compute chunked block size from sliding window if needed.
1132
+ chunk_size = -1
1133
+ if sliding_window_val > 0 and chunk_lookback > -1:
1134
+ chunk_size = sliding_window_val // (chunk_lookback + 1)
1135
+ assert chunk_size > 0, "sliding_window must be > chunk_lookback+1"
1136
+ elif sliding_window_val <= 0:
1137
+ chunk_lookback = -1
1138
+
1139
TILE_SIZE_PREFILL = _get_tile_size(
1140
head_size,
1141
sliding_window_val,
@@ -1184,6 +1227,8 @@ def unified_attention(
1227
stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0,
1228
stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0,
1229
stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0,
1230
+ CHUNK_LOOKBACK=chunk_lookback,
1231
+ CHUNK_SIZE=chunk_size,
1232
)
1233
else:
1234
kernel_unified_attention_3d[
@@ -1245,6 +1290,8 @@ def unified_attention(
1290
stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0,
1291
stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0,
1292
stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0,
1293
+ CHUNK_LOOKBACK=chunk_lookback,
1294
+ CHUNK_SIZE=chunk_size,
1295
)
1296
reduce_segments[(q.shape[0], num_query_heads)](
1297
output_ptr=out,