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, 07:45 UTC
[Bugfix] Fix missing scale passing for encoder Triton Attention implementation (#32149)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jan 12, 2026 at 19:13 UTC
9dbe1fe960f9931a262c88b4082aa7551cf7dfc0
4 files changed
+13
-27
examples/offline_inference/basic/embed.py
-8
@@ -4,10 +4,7 @@
4
from argparse import Namespace
5
6
from vllm import LLM, EngineArgs
7
-from vllm.config import AttentionConfig
8
-from vllm.platforms import current_platform
7
from vllm.utils.argparse_utils import FlexibleArgumentParser
10
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
8
9
10
def parse_args():
@@ -23,11 +20,6 @@ def parse_args():
20
21
22
def main(args: Namespace):
26
- if current_platform.is_rocm():
27
- args.attention_config = AttentionConfig(
28
- backend=AttentionBackendEnum.FLEX_ATTENTION
29
- )
30
-
23
# Sample prompts.
24
prompts = [
25
"Hello, my name is",
examples/offline_inference/basic/score.py
-8
@@ -4,10 +4,7 @@
4
from argparse import Namespace
5
6
from vllm import LLM, EngineArgs
7
-from vllm.config import AttentionConfig
8
-from vllm.platforms import current_platform
7
from vllm.utils.argparse_utils import FlexibleArgumentParser
10
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
8
9
10
def parse_args():
@@ -23,11 +20,6 @@ def parse_args():
20
21
22
def main(args: Namespace):
26
- if current_platform.is_rocm():
27
- args.attention_config = AttentionConfig(
28
- backend=AttentionBackendEnum.FLEX_ATTENTION
29
- )
30
-
23
# Sample prompts.
24
text_1 = "What is the capital of France?"
25
texts_2 = [
vllm/v1/attention/backends/triton_attn.py
+1
@@ -573,6 +573,7 @@ class TritonAttentionImpl(AttentionImpl):
573
b_seq_len=seq_lens,
574
max_input_len=max_query_len,
575
is_causal=False, # Encoder attention is bidirectional
576
+ softmax_scale=self.scale,
577
sliding_window_q=self.sliding_window[0],
578
sliding_window_k=self.sliding_window[1],
579
)
vllm/v1/attention/ops/triton_prefill_attention.py
+12
-11
@@ -211,16 +211,17 @@ def get_block_size(dtype: torch.dtype) -> int:
211
212
213
def context_attention_fwd(
214
- q,
215
- k,
216
- v,
217
- o,
218
- b_start_loc,
219
- b_seq_len,
220
- max_input_len,
221
- is_causal=True,
222
- sliding_window_q=None,
223
- sliding_window_k=None,
214
+ q: torch.Tensor,
215
+ k: torch.Tensor,
216
+ v: torch.Tensor,
217
+ o: torch.Tensor,
218
+ b_start_loc: torch.Tensor,
219
+ b_seq_len: torch.Tensor,
220
+ max_input_len: int,
221
+ is_causal: bool = True,
222
+ softmax_scale: float | None = None,
223
+ sliding_window_q: int | None = None,
224
+ sliding_window_k: int | None = None,
225
):
226
"""
227
q, k, v: [b * s, head, head_dim]
@@ -232,7 +233,7 @@ def context_attention_fwd(
233
234
Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1]
235
235
- sm_scale = 1.0 / (Lq**0.5)
236
+ sm_scale = 1.0 / (Lq**0.5) if softmax_scale is None else softmax_scale
237
batch, head = b_seq_len.shape[0], q.shape[1]
238
kv_group_num = q.shape[1] // k.shape[1]
239