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 10, 01:30 UTC
[Model] Support Step1 Model (#32511)
Signed-off-by: xieli <xieli@stepfun.com>
Li Xie committed
Jan 18, 2026 at 18:20 UTC
c826c72a9633454679871fcb81fbc31fe03fb150
9 files changed
+472
-6
docs/models/supported_models.md
+2
-1
@@ -452,9 +452,10 @@ th {
452
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
453
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
454
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
455
+| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
456
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
457
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
457
-| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
458
+| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
459
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
460
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
461
| `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+3
@@ -472,6 +472,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
472
"ByteDance-Seed/Seed-OSS-36B-Instruct",
473
trust_remote_code=True,
474
),
475
+ "Step1ForCausalLM": _HfExamplesInfo(
476
+ "stepfun-ai/Step-Audio-EditX", trust_remote_code=True
477
+ ),
478
"SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
479
"StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),
480
"StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
tests/models/test_initialization.py
+4
-1
@@ -115,8 +115,11 @@ def can_initialize(
115
# FIXME: A hack to bypass FA3 assertion because our CI's L4 GPU
116
# has cc==8.9 which hasn't supported FA3 yet. Remove this hack when
117
# L4 supports FA3.
118
+ # Step1ForCausalLM requires TRITON_ATTN for use_alibi_sqrt support.
119
attention_config = (
119
- {"backend": "TRITON_ATTN"} if model_arch == "GptOssForCausalLM" else None
120
+ {"backend": "TRITON_ATTN"}
121
+ if model_arch in ("GptOssForCausalLM", "Step1ForCausalLM")
122
+ else None
123
)
124
if model_arch == "WhisperForConditionalGeneration":
125
m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
vllm/attention/layer.py
+11
-1
@@ -162,6 +162,7 @@ class Attention(nn.Module, AttentionLayerBase):
162
scale: float,
163
num_kv_heads: int | None = None,
164
alibi_slopes: list[float] | None = None,
165
+ use_alibi_sqrt: bool | None = None,
166
cache_config: CacheConfig | None = None,
167
quant_config: QuantizationConfig | None = None,
168
logits_soft_cap: float | None = None,
@@ -243,7 +244,16 @@ class Attention(nn.Module, AttentionLayerBase):
244
)
245
else:
246
self.attn_backend = attn_backend
246
-
247
+ backend_supports_alibi_sqrt = self.attn_backend.supports_alibi_sqrt()
248
+ use_alibi_sqrt = use_alibi_sqrt if use_alibi_sqrt else False
249
+ if use_alibi_sqrt and not backend_supports_alibi_sqrt:
250
+ raise ValueError(
251
+ f"use_alibi_sqrt is not supported by backend "
252
+ f"{self.attn_backend.get_name()}."
253
+ )
254
+ self.use_alibi_sqrt = bool(use_alibi_sqrt)
255
+ if backend_supports_alibi_sqrt:
256
+ extra_impl_args["use_alibi_sqrt"] = self.use_alibi_sqrt
257
# prefix caching + batch invariance is currently not supported for
258
# FLASHINFER and TRITON_MLA.
259
if (
vllm/model_executor/models/registry.py
+1
@@ -185,6 +185,7 @@ _TEXT_GENERATION_MODELS = {
185
"Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
186
"RWForCausalLM": ("falcon", "FalconForCausalLM"),
187
"SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
188
+ "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
189
"Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
190
"StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
191
"StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
vllm/model_executor/models/step1.py
new
+415
@@ -0,0 +1,415 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared Step decoder blocks and the Step1 text model."""
4
+
5
+from __future__ import annotations
6
+
7
+import math
8
+from collections.abc import Iterable
9
+
10
+import torch
11
+from torch import nn
12
+
13
+from vllm.attention.layer import Attention, AttentionType
14
+from vllm.config import CacheConfig, VllmConfig
15
+from vllm.distributed import (
16
+ get_pp_group,
17
+ get_tensor_model_parallel_rank,
18
+ get_tensor_model_parallel_world_size,
19
+)
20
+from vllm.model_executor.layers.activation import SiluAndMul
21
+from vllm.model_executor.layers.layernorm import RMSNorm
22
+from vllm.model_executor.layers.linear import (
23
+ MergedColumnParallelLinear,
24
+ QKVParallelLinear,
25
+ RowParallelLinear,
26
+)
27
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
28
+from vllm.model_executor.layers.quantization import QuantizationConfig
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 default_weight_loader
34
+from vllm.model_executor.models.interfaces import SupportsPP
35
+from vllm.model_executor.models.utils import (
36
+ AutoWeightsLoader,
37
+ PPMissingLayer,
38
+ is_pp_missing_parameter,
39
+ make_empty_intermediate_tensors_factory,
40
+ make_layers,
41
+ maybe_prefix,
42
+)
43
+from vllm.sequence import IntermediateTensors
44
+
45
+STEP_PACKED_MODULES_MAPPING = {
46
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
47
+ "gate_up_proj": ["gate_proj", "up_proj"],
48
+}
49
+
50
+
51
+def _get_step_alibi_slopes(total_num_heads: int) -> torch.Tensor:
52
+ """Reference ALiBi slopes used by Step models."""
53
+ closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
54
+ base = torch.tensor(
55
+ 2 ** (-8.0 / closest_power_of_2),
56
+ dtype=torch.float32,
57
+ )
58
+ slopes = torch.pow(
59
+ base,
60
+ torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32),
61
+ )
62
+ if closest_power_of_2 != total_num_heads:
63
+ extra_base = torch.tensor(
64
+ 2 ** (-4.0 / closest_power_of_2),
65
+ dtype=torch.float32,
66
+ )
67
+ num_remaining_heads = total_num_heads - closest_power_of_2
68
+ extra_powers = torch.arange(
69
+ 1,
70
+ 1 + 2 * num_remaining_heads,
71
+ 2,
72
+ dtype=torch.int32,
73
+ )
74
+ slopes = torch.cat(
75
+ [slopes, torch.pow(extra_base, extra_powers)],
76
+ dim=0,
77
+ )
78
+ return slopes
79
+
80
+
81
+class StepAttention(nn.Module):
82
+ def __init__(
83
+ self,
84
+ config,
85
+ cache_config: CacheConfig | None = None,
86
+ quant_config: QuantizationConfig | None = None,
87
+ prefix: str = "",
88
+ ):
89
+ super().__init__()
90
+ self.hidden_size = config.hidden_size
91
+ tp_size = get_tensor_model_parallel_world_size()
92
+ self.total_num_heads = config.num_attention_heads
93
+ assert self.total_num_heads % tp_size == 0
94
+ self.num_heads = self.total_num_heads // tp_size
95
+ self.head_dim = self.hidden_size // self.total_num_heads
96
+
97
+ total_num_kv_heads = getattr(
98
+ config, "num_attention_groups", getattr(config, "num_key_value_heads", 1)
99
+ )
100
+ if total_num_kv_heads is None or total_num_kv_heads <= 0:
101
+ total_num_kv_heads = 1
102
+ self.total_num_kv_heads = total_num_kv_heads
103
+ if self.total_num_kv_heads >= tp_size:
104
+ assert self.total_num_kv_heads % tp_size == 0
105
+ else:
106
+ assert tp_size % self.total_num_kv_heads == 0
107
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
108
+
109
+ self.qkv_proj = QKVParallelLinear(
110
+ hidden_size=self.hidden_size,
111
+ head_size=self.head_dim,
112
+ total_num_heads=self.total_num_heads,
113
+ total_num_kv_heads=self.total_num_kv_heads,
114
+ bias=getattr(config, "attention_bias", False),
115
+ quant_config=quant_config,
116
+ prefix=f"{prefix}.qkv_proj",
117
+ )
118
+ self.q_size = self.num_heads * self.head_dim
119
+ self.kv_size = self.num_kv_heads * self.head_dim
120
+
121
+ self.o_proj = RowParallelLinear(
122
+ input_size=self.total_num_heads * self.head_dim,
123
+ output_size=self.hidden_size,
124
+ bias=getattr(config, "attention_bias", False),
125
+ quant_config=quant_config,
126
+ prefix=f"{prefix}.o_proj",
127
+ )
128
+
129
+ tp_rank = get_tensor_model_parallel_rank()
130
+ head_start = tp_rank * self.num_heads
131
+ head_end = (tp_rank + 1) * self.num_heads
132
+ alibi_slopes = _get_step_alibi_slopes(self.total_num_heads)[head_start:head_end]
133
+ alibi_slopes = alibi_slopes.tolist()
134
+
135
+ self.scale = self.head_dim**-0.5
136
+ self.attn = Attention(
137
+ self.num_heads,
138
+ self.head_dim,
139
+ self.scale,
140
+ num_kv_heads=self.num_kv_heads,
141
+ cache_config=cache_config,
142
+ quant_config=quant_config,
143
+ alibi_slopes=alibi_slopes,
144
+ prefix=f"{prefix}.attn",
145
+ use_alibi_sqrt=True,
146
+ attn_type=AttentionType.DECODER,
147
+ )
148
+
149
+ def forward(
150
+ self,
151
+ hidden_states: torch.Tensor,
152
+ ) -> torch.Tensor:
153
+ qkv, _ = self.qkv_proj(hidden_states)
154
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
155
+ attn_output = self.attn(q, k, v)
156
+ output, _ = self.o_proj(attn_output)
157
+ return output
158
+
159
+
160
+class StepMLP(nn.Module):
161
+ def __init__(
162
+ self,
163
+ hidden_size: int,
164
+ intermediate_size: int,
165
+ quant_config: QuantizationConfig | None = None,
166
+ prefix: str = "",
167
+ bias: bool = False,
168
+ ):
169
+ super().__init__()
170
+ self.gate_up_proj = MergedColumnParallelLinear(
171
+ input_size=hidden_size,
172
+ output_sizes=[intermediate_size, intermediate_size],
173
+ bias=bias,
174
+ quant_config=quant_config,
175
+ prefix=f"{prefix}.gate_up_proj",
176
+ )
177
+ self.down_proj = RowParallelLinear(
178
+ input_size=intermediate_size,
179
+ output_size=hidden_size,
180
+ bias=bias,
181
+ quant_config=quant_config,
182
+ prefix=f"{prefix}.down_proj",
183
+ )
184
+ self.act_fn = SiluAndMul()
185
+
186
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
187
+ x, _ = self.gate_up_proj(x)
188
+ x = self.act_fn(x)
189
+ x, _ = self.down_proj(x)
190
+ return x
191
+
192
+
193
+class StepDecoderLayer(nn.Module):
194
+ def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
195
+ super().__init__()
196
+ config = vllm_config.model_config.hf_config
197
+ cache_config = vllm_config.cache_config
198
+ quant_config = vllm_config.quant_config
199
+
200
+ self.hidden_size = config.hidden_size
201
+ self.self_attn = StepAttention(
202
+ config=config,
203
+ cache_config=cache_config,
204
+ quant_config=quant_config,
205
+ prefix=f"{prefix}.self_attn",
206
+ )
207
+ self.mlp = StepMLP(
208
+ hidden_size=self.hidden_size,
209
+ intermediate_size=config.intermediate_size,
210
+ quant_config=quant_config,
211
+ prefix=f"{prefix}.mlp",
212
+ bias=getattr(config, "mlp_bias", False),
213
+ )
214
+ self.input_layernorm = RMSNorm(
215
+ self.hidden_size,
216
+ eps=config.rms_norm_eps,
217
+ )
218
+ self.post_attention_layernorm = RMSNorm(
219
+ self.hidden_size,
220
+ eps=config.rms_norm_eps,
221
+ )
222
+
223
+ def forward(
224
+ self,
225
+ positions: torch.Tensor,
226
+ hidden_states: torch.Tensor,
227
+ residual: torch.Tensor | None,
228
+ ) -> tuple[torch.Tensor, torch.Tensor]:
229
+ if residual is None:
230
+ residual = hidden_states
231
+ hidden_states = self.input_layernorm(hidden_states)
232
+ else:
233
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
234
+ hidden_states = self.self_attn(hidden_states=hidden_states)
235
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
236
+ hidden_states = self.mlp(hidden_states)
237
+ return hidden_states, residual
238
+
239
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
240
+ stacked_params_mapping = [
241
+ (".qkv_proj", ".q_proj", "q"),
242
+ (".qkv_proj", ".k_proj", "k"),
243
+ (".qkv_proj", ".v_proj", "v"),
244
+ (".gate_up_proj", ".gate_proj", 0),
245
+ (".gate_up_proj", ".up_proj", 1),
246
+ ]
247
+ params_dict = dict(self.named_parameters())
248
+ loaded_params: set[str] = set()
249
+ for name, loaded_weight in weights:
250
+ for param_name, weight_name, shard_id in stacked_params_mapping:
251
+ if weight_name not in name:
252
+ continue
253
+ name = name.replace(weight_name, param_name)
254
+ # Skip loading extra bias for GPTQ models.
255
+ if name.endswith(".bias") and name not in params_dict:
256
+ continue
257
+ if is_pp_missing_parameter(name, self):
258
+ continue
259
+ param = params_dict[name]
260
+ weight_loader = param.weight_loader
261
+ weight_loader(param, loaded_weight, shard_id)
262
+ break
263
+ else:
264
+ # Skip loading extra bias for GPTQ models.
265
+ if name.endswith(".bias") and name not in params_dict:
266
+ continue
267
+ if is_pp_missing_parameter(name, self):
268
+ continue
269
+ param = params_dict[name]
270
+ weight_loader = getattr(param, "weight_loader", default_weight_loader) # type: ignore[name-defined]
271
+ weight_loader(param, loaded_weight)
272
+ loaded_params.add(name)
273
+ return loaded_params
274
+
275
+
276
+class StepDecoderModel(nn.Module):
277
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
278
+ super().__init__()
279
+ config = vllm_config.model_config.hf_config
280
+ quant_config = vllm_config.quant_config
281
+
282
+ self.config = config
283
+ self.quant_config = quant_config
284
+ # Need embed_tokens on first rank, and also on last rank if tie_word_embeddings
285
+ if get_pp_group().is_first_rank or (
286
+ config.tie_word_embeddings and get_pp_group().is_last_rank
287
+ ):
288
+ self.embed_tokens = VocabParallelEmbedding(
289
+ config.vocab_size,
290
+ config.hidden_size,
291
+ quant_config=quant_config,
292
+ )
293
+ else:
294
+ self.embed_tokens = PPMissingLayer()
295
+ self.start_layer, self.end_layer, self.layers = make_layers(
296
+ config.num_hidden_layers,
297
+ lambda prefix: StepDecoderLayer(vllm_config=vllm_config, prefix=prefix),
298
+ prefix=maybe_prefix(prefix, "layers"),
299
+ )
300
+ if get_pp_group().is_last_rank:
301
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
302
+ else:
303
+ self.norm = PPMissingLayer()
304
+
305
+ self.aux_hidden_state_layers: tuple[int, ...] = getattr(
306
+ config, "aux_hidden_state_layers", ()
307
+ )
308
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
309
+ ["hidden_states", "residual"],
310
+ config.hidden_size,
311
+ )
312
+
313
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
314
+ return self.embed_tokens(input_ids)
315
+
316
+ def forward(
317
+ self,
318
+ input_ids: torch.Tensor | None,
319
+ positions: torch.Tensor,
320
+ intermediate_tensors: IntermediateTensors | None,
321
+ inputs_embeds: torch.Tensor | None = None,
322
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
323
+ if get_pp_group().is_first_rank:
324
+ if inputs_embeds is not None:
325
+ hidden_states = inputs_embeds
326
+ else:
327
+ assert input_ids is not None
328
+ hidden_states = self.embed_input_ids(input_ids)
329
+ residual = None
330
+ else:
331
+ assert intermediate_tensors is not None
332
+ hidden_states = intermediate_tensors["hidden_states"]
333
+ residual = intermediate_tensors["residual"]
334
+
335
+ aux_hidden_states = []
336
+ for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
337
+ if idx in self.aux_hidden_state_layers:
338
+ if residual is None:
339
+ aux_hidden_states.append(hidden_states)
340
+ else:
341
+ aux_hidden_states.append(hidden_states + residual)
342
+ hidden_states, residual = layer(positions, hidden_states, residual)
343
+
344
+ if not get_pp_group().is_last_rank:
345
+ return IntermediateTensors(
346
+ {"hidden_states": hidden_states, "residual": residual}
347
+ )
348
+
349
+ hidden_states, _ = self.norm(hidden_states, residual)
350
+ if aux_hidden_states:
351
+ return hidden_states, aux_hidden_states
352
+ return hidden_states
353
+
354
+
355
+class Step1ForCausalLM(nn.Module, SupportsPP):
356
+ packed_modules_mapping = STEP_PACKED_MODULES_MAPPING
357
+
358
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
359
+ super().__init__()
360
+ config = vllm_config.model_config.hf_config
361
+ quant_config = vllm_config.quant_config
362
+
363
+ self.config = config
364
+ self.quant_config = quant_config
365
+ self.model = StepDecoderModel(
366
+ vllm_config=vllm_config,
367
+ prefix=maybe_prefix(prefix, "model"),
368
+ )
369
+
370
+ if get_pp_group().is_last_rank:
371
+ self.lm_head = ParallelLMHead(
372
+ config.vocab_size,
373
+ config.hidden_size,
374
+ quant_config=quant_config,
375
+ prefix=maybe_prefix(prefix, "lm_head"),
376
+ )
377
+ if getattr(config, "tie_word_embeddings", True):
378
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
379
+ self.logits_processor = LogitsProcessor(config.vocab_size)
380
+ else:
381
+ self.lm_head = PPMissingLayer()
382
+ self.logits_processor = None # type: ignore[assignment]
383
+
384
+ self.make_empty_intermediate_tensors = (
385
+ self.model.make_empty_intermediate_tensors
386
+ )
387
+
388
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
389
+ return self.model.embed_input_ids(input_ids)
390
+
391
+ def forward(
392
+ self,
393
+ input_ids: torch.LongTensor | None,
394
+ positions: torch.Tensor,
395
+ intermediate_tensors: IntermediateTensors | None,
396
+ inputs_embeds: torch.Tensor | None = None,
397
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
398
+ return self.model(
399
+ input_ids,
400
+ positions,
401
+ intermediate_tensors,
402
+ inputs_embeds=inputs_embeds,
403
+ )
404
+
405
+ def compute_logits(
406
+ self,
407
+ hidden_states: torch.Tensor,
408
+ ) -> torch.Tensor | None:
409
+ if not get_pp_group().is_last_rank:
410
+ return None
411
+ return self.logits_processor(self.lm_head, hidden_states)
412
+
413
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
414
+ loader = AutoWeightsLoader(self)
415
+ return loader.load_weights(weights)
vllm/v1/attention/backend.py
+4
@@ -172,6 +172,10 @@ class AttentionBackend(ABC):
172
def supports_sink(cls) -> bool:
173
return False
174
175
+ @classmethod
176
+ def supports_alibi_sqrt(cls) -> bool:
177
+ return False
178
+
179
@classmethod
180
def supports_mm_prefix(cls) -> bool:
181
return False
vllm/v1/attention/backends/triton_attn.py
+7
-1
@@ -331,6 +331,10 @@ class TritonAttentionBackend(AttentionBackend):
331
AttentionType.ENCODER_DECODER,
332
)
333
334
+ @classmethod
335
+ def supports_alibi_sqrt(cls) -> bool:
336
+ return True
337
+
338
@classmethod
339
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
340
return True
@@ -353,6 +357,7 @@ class TritonAttentionImpl(AttentionImpl):
357
attn_type: AttentionType = AttentionType.DECODER,
358
kv_sharing_target_layer_name: int | None = None,
359
sinks: torch.Tensor | None = None,
360
+ use_alibi_sqrt: bool = False,
361
) -> None:
362
self.num_heads = num_heads
363
self.head_size = head_size
@@ -386,7 +391,7 @@ class TritonAttentionImpl(AttentionImpl):
391
f"heads in the layer. Sinks shape: {sinks.shape}, "
392
f"num_heads: {num_heads}."
393
)
389
-
394
+ self.use_alibi_sqrt = use_alibi_sqrt
395
self.supports_quant_query_input = current_platform.is_cuda()
396
397
def forward(
@@ -513,6 +518,7 @@ class TritonAttentionImpl(AttentionImpl):
518
softmax_scale=self.scale,
519
causal=True,
520
alibi_slopes=self.alibi_slopes,
521
+ use_alibi_sqrt=self.use_alibi_sqrt,
522
window_size=self.sliding_window,
523
block_table=block_table,
524
softcap=self.logits_soft_cap,
vllm/v1/attention/ops/triton_unified_attention.py
+25
-2
@@ -82,6 +82,7 @@ def kernel_unified_attention_2d(
82
HEAD_SIZE: tl.constexpr, # int
83
HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
84
USE_ALIBI_SLOPES: tl.constexpr, # bool
85
+ USE_ALIBI_SQRT: tl.constexpr, # bool
86
USE_QQ_BIAS: tl.constexpr, # bool
87
USE_SOFTCAP: tl.constexpr, # bool
88
USE_SINKS: tl.constexpr, # bool
@@ -325,7 +326,16 @@ def kernel_unified_attention_2d(
326
)
327
328
if USE_ALIBI_SLOPES:
328
- S += alibi_slope[:, None] * (seq_offset - context_len)
329
+ if USE_ALIBI_SQRT:
330
+ relative_pos = seq_offset - (context_len + query_pos[:, None])
331
+ alibi_offset = tl.where(
332
+ relative_pos <= 0,
333
+ -tl.sqrt((-relative_pos).to(tl.float32)),
334
+ 0.0,
335
+ )
336
+ else:
337
+ alibi_offset = seq_offset - context_len
338
+ S += alibi_slope[:, None] * alibi_offset
339
340
if USE_QQ_BIAS:
341
# compute key positions relative to query section
@@ -420,6 +430,7 @@ def kernel_unified_attention_3d(
430
HEAD_SIZE: tl.constexpr, # int
431
HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
432
USE_ALIBI_SLOPES: tl.constexpr, # bool
433
+ USE_ALIBI_SQRT: tl.constexpr, # bool
434
USE_QQ_BIAS: tl.constexpr, # bool
435
USE_SOFTCAP: tl.constexpr, # bool
436
USE_SINKS: tl.constexpr, # bool
@@ -669,7 +680,16 @@ def kernel_unified_attention_3d(
680
)
681
682
if USE_ALIBI_SLOPES:
672
- S += alibi_slope[:, None] * (seq_offset - context_len)
683
+ if USE_ALIBI_SQRT:
684
+ relative_pos = seq_offset - (context_len + query_pos[:, None])
685
+ alibi_offset = tl.where(
686
+ relative_pos <= 0,
687
+ -tl.sqrt((-relative_pos).to(tl.float32)),
688
+ 0.0,
689
+ )
690
+ else:
691
+ alibi_offset = seq_offset - context_len
692
+ S += alibi_slope[:, None] * alibi_offset
693
694
if USE_QQ_BIAS:
695
# compute key positions relative to query section
@@ -888,6 +908,7 @@ def unified_attention(
908
sinks=None,
909
# Optional tensor for prefix lengths (PrefixLM support)
910
mm_prefix_range=None,
911
+ use_alibi_sqrt=False,
912
):
913
assert causal, "Only causal attention is supported"
914
assert q_descale is None, "Q scales not supported"
@@ -994,6 +1015,7 @@ def unified_attention(
1015
HEAD_SIZE=head_size,
1016
HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
1017
USE_ALIBI_SLOPES=use_alibi_slopes,
1018
+ USE_ALIBI_SQRT=use_alibi_sqrt,
1019
USE_QQ_BIAS=use_qq_bias,
1020
USE_SOFTCAP=(softcap > 0),
1021
USE_SINKS=(sinks is not None),
@@ -1045,6 +1067,7 @@ def unified_attention(
1067
HEAD_SIZE=head_size,
1068
HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
1069
USE_ALIBI_SLOPES=use_alibi_slopes,
1070
+ USE_ALIBI_SQRT=use_alibi_sqrt,
1071
USE_QQ_BIAS=use_qq_bias,
1072
USE_SOFTCAP=(softcap > 0),
1073
USE_SINKS=(sinks is not None),