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 15, 19:45 UTC
[FEAT] [ROCm] [Embedding] Add encoder-only model support into ROCm Flash Attention to enable embedding models. (#14664)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Mar 13, 2025 at 00:31 UTC
916836bbfb7ef52077e78602439a944298dbf886
7 files changed
+118
-50
CMakeLists.txt
+4
@@ -561,6 +561,10 @@ set(VLLM_MOE_EXT_SRC
561
"csrc/moe/moe_align_sum_kernels.cu"
562
"csrc/moe/topk_softmax_kernels.cu")
563
564
+if(VLLM_GPU_LANG STREQUAL "CUDA")
565
+ list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu")
566
+endif()
567
+
568
set_gencode_flags_for_srcs(
569
SRCS "${VLLM_MOE_EXT_SRC}"
570
CUDA_ARCHS "${CUDA_ARCHS}")
csrc/moe/torch_bindings.cpp
+1
@@ -52,6 +52,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
52
"int moe_block_size, bool replicate_input, bool apply_weights)"
53
" -> Tensor");
54
// conditionally compiled so impl registration is in source file
55
+
56
#endif
57
}
58
tests/models/embedding/language/test_cls_models.py
+15
-2
@@ -7,6 +7,8 @@ import pytest
7
import torch
8
from transformers import AutoModelForSequenceClassification
9
10
+from vllm.platforms import current_platform
11
+
12
13
@pytest.mark.parametrize(
14
"model",
@@ -15,14 +17,21 @@ from transformers import AutoModelForSequenceClassification
17
marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
18
],
19
)
18
-@pytest.mark.parametrize("dtype", ["float"])
20
+@pytest.mark.parametrize("dtype",
21
+ ["half"] if current_platform.is_rocm() else ["float"])
22
def test_classification_models(
23
hf_runner,
24
vllm_runner,
25
example_prompts,
26
model: str,
27
dtype: str,
28
+ monkeypatch,
29
) -> None:
30
+ if current_platform.is_rocm():
31
+ # ROCm Triton FA does not currently support sliding window attention
32
+ # switch to use ROCm CK FA backend
33
+ monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")
34
+
35
with vllm_runner(model, dtype=dtype) as vllm_model:
36
vllm_outputs = vllm_model.classify(example_prompts)
37
@@ -43,4 +52,8 @@ def test_classification_models(
52
hf_output = torch.tensor(hf_output)
53
vllm_output = torch.tensor(vllm_output)
54
46
- assert torch.allclose(hf_output, vllm_output, 1e-3)
55
+ # the tolerance value of 1e-2 is selected based on the
56
+ # half datatype tests in
57
+ # tests/models/embedding/language/test_embedding.py
58
+ assert torch.allclose(hf_output, vllm_output,
59
+ 1e-3 if dtype == "float" else 1e-2)
tests/models/embedding/language/test_embedding.py
+11
-2
@@ -6,6 +6,7 @@ Run `pytest tests/models/embedding/language/test_embedding.py`.
6
import pytest
7
8
from vllm.config import PoolerConfig
9
+from vllm.platforms import current_platform
10
11
from ..utils import check_embeddings_close
12
@@ -18,15 +19,15 @@ from ..utils import check_embeddings_close
19
marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
20
pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
21
pytest.param("intfloat/multilingual-e5-small"),
22
+ pytest.param("Alibaba-NLP/gte-Qwen2-7B-instruct"),
23
# [Decoder-only]
24
pytest.param("BAAI/bge-multilingual-gemma2",
25
marks=[pytest.mark.core_model]),
26
pytest.param("intfloat/e5-mistral-7b-instruct",
27
marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
28
pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct"),
27
- pytest.param("Alibaba-NLP/gte-Qwen2-7B-instruct"),
29
pytest.param("ssmits/Qwen2-7B-Instruct-embed-base"),
29
- # [Encoder-decoder]
30
+ # [Cross-Encoder]
31
pytest.param("sentence-transformers/stsb-roberta-base-v2"),
32
],
33
)
@@ -37,11 +38,19 @@ def test_models(
38
example_prompts,
39
model,
40
dtype: str,
41
+ monkeypatch,
42
) -> None:
43
+
44
+ if model == "BAAI/bge-multilingual-gemma2" and current_platform.is_rocm():
45
+ # ROCm Triton FA does not currently support sliding window attention
46
+ # switch to use ROCm CK FA backend
47
+ monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")
48
+
49
vllm_extra_kwargs = {}
50
if model == "ssmits/Qwen2-7B-Instruct-embed-base":
51
vllm_extra_kwargs["override_pooler_config"] = \
52
PoolerConfig(pooling_type="MEAN")
53
+
54
if model == "Alibaba-NLP/gte-Qwen2-1.5B-instruct":
55
vllm_extra_kwargs["hf_overrides"] = {"is_causal": True}
56
tests/models/embedding/language/test_gritlm.py
+2
-2
@@ -15,8 +15,8 @@ import vllm.config
15
from ....utils import RemoteOpenAIServer
16
17
# GritLM embedding implementation is only supported by XFormers backend.
18
-pytest.mark.skipif(not importlib.util.find_spec("xformers"),
19
- reason="GritLM requires XFormers")
18
+pytestmark = pytest.mark.skipif(not importlib.util.find_spec("xformers"),
19
+ reason="GritLM requires XFormers")
20
21
MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
22
MAX_MODEL_LEN = 4000
tests/models/embedding/vision_language/test_llava_next.py
+17
@@ -4,10 +4,27 @@ import pytest
4
import torch.nn.functional as F
5
from transformers import AutoModelForVision2Seq
6
7
+from vllm.platforms import current_platform
8
+
9
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
10
from ....utils import large_gpu_test
11
from ..utils import check_embeddings_close
12
13
+# Llava Next embedding implementation is only supported by CUDA.
14
+# If run on ROCm, hf_model.model.resize_token_embeddings will
15
+# cause the following error:
16
+# RuntimeError: Calling torch.linalg.cholesky on a CUDA tensor
17
+# requires compiling PyTorch with MAGMA. Please use PyTorch
18
+# built with MAGMA support.
19
+# If run on CPU, hf_model.model.resize_token_embeddings will
20
+# cause the following error:
21
+# RuntimeError: Calling torch.linalg.cholesky on a CPU tensor
22
+# requires compiling PyTorch with LAPACK. Please use PyTorch
23
+# built with LAPACK support.
24
+pytestmark = pytest.mark.skipif(
25
+ not current_platform.is_cuda(),
26
+ reason="Llava Next model uses op that is only supported in CUDA")
27
+
28
llama3_template = '<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n' # noqa: E501
29
30
HF_TEXT_PROMPTS = [
vllm/attention/backends/rocm_flash_attn.py
+68
-44
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""Attention layer ROCm GPUs."""
3
+import itertools
4
from dataclasses import dataclass
5
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
6
@@ -342,28 +343,27 @@ def _get_seq_len_block_table_args(
343
Decoder attn -> select entirely decoder self-attention-related fields
344
Encoder/decoder cross-attn -> select encoder sequence lengths
345
Encoder attn -> select encoder sequence lengths fields
346
+ Encoder-only attn -> select prefill sequence lengths with
347
+ bidirectional attention
348
349
Arguments:
350
351
* attn_metadata: Attention metadata structure associated with attention op
352
* attn_type: encoder attention, decoder self-attention,
350
- encoder/decoder cross-attention
353
+ encoder/decoder cross-attention, encoder-only
354
355
Returns:
356
357
* Appropriate sequence-lengths tensors for query and key
358
* Appropriate max sequence-length scalar
359
+ * Causal masking flag
360
'''
361
358
- partial_prefix_sum = 0
362
if attn_type == AttentionType.ENCODER:
363
assert attn_metadata.encoder_seq_lens is not None
364
assert attn_metadata.encoder_seq_lens_tensor is not None
365
query_seq_start_loc = torch.tensor(
363
- [0] + [
364
- partial_prefix_sum := partial_prefix_sum + i
365
- for i in attn_metadata.encoder_seq_lens
366
- ],
366
+ list(itertools.accumulate([0] + attn_metadata.encoder_seq_lens)),
367
device=attn_metadata.encoder_seq_lens_tensor.device,
368
dtype=attn_metadata.encoder_seq_lens_tensor.dtype)
369
causal_mask = False
@@ -372,16 +372,29 @@ def _get_seq_len_block_table_args(
372
return (query_seq_start_loc, attn_metadata.max_encoder_seq_len,
373
query_seq_start_loc, attn_metadata.max_encoder_seq_len,
374
attn_metadata.encoder_seq_lens, causal_mask)
375
+
376
+ elif attn_type == AttentionType.ENCODER_ONLY:
377
+ # For encoder-only models, we use the prefill sequence lengths
378
+ assert attn_metadata.seq_lens is not None
379
+ assert attn_metadata.seq_lens_tensor is not None
380
+ query_seq_start_loc = torch.tensor(
381
+ list(itertools.accumulate([0] + attn_metadata.seq_lens)),
382
+ device=attn_metadata.seq_lens_tensor.device,
383
+ dtype=attn_metadata.seq_lens_tensor.dtype)
384
+ max_seq_len = attn_metadata.max_prefill_seq_len
385
+ # Encoder-only models typically use bidirectional attention
386
+ causal_mask = False
387
+
388
+ return (query_seq_start_loc, max_seq_len, query_seq_start_loc,
389
+ max_seq_len, attn_metadata.seq_lens, causal_mask)
390
+
391
elif attn_type == AttentionType.DECODER:
392
# Decoder self-attention
393
# Choose max_seq_len based on whether we are in prompt_run
394
assert attn_metadata.seq_lens is not None
395
assert attn_metadata.seq_lens_tensor is not None
396
query_seq_start_loc = torch.tensor(
381
- [0] + [
382
- partial_prefix_sum := partial_prefix_sum + i
383
- for i in attn_metadata.seq_lens
384
- ],
397
+ list(itertools.accumulate([0] + attn_metadata.seq_lens)),
398
device=attn_metadata.seq_lens_tensor.device,
399
dtype=attn_metadata.seq_lens_tensor.dtype)
400
max_seq_len = attn_metadata.max_prefill_seq_len
@@ -393,21 +406,14 @@ def _get_seq_len_block_table_args(
406
assert attn_metadata.seq_lens is not None
407
assert attn_metadata.encoder_seq_lens_tensor is not None
408
query_start_loc = torch.tensor(
396
- [0] + [
397
- partial_prefix_sum := partial_prefix_sum + i
398
- for i in attn_metadata.seq_lens
399
- ],
409
+ list(itertools.accumulate([0] + attn_metadata.seq_lens)),
410
device=attn_metadata.encoder_seq_lens_tensor.device,
411
dtype=attn_metadata.encoder_seq_lens_tensor.dtype)
412
403
- partial_prefix_sum = 0
413
assert attn_metadata.encoder_seq_lens is not None
414
assert attn_metadata.seq_lens_tensor is not None
415
key_seq_start_loc = torch.tensor(
407
- [0] + [
408
- partial_prefix_sum := partial_prefix_sum + i
409
- for i in attn_metadata.encoder_seq_lens
410
- ],
416
+ list(itertools.accumulate([0] + attn_metadata.encoder_seq_lens)),
417
device=attn_metadata.seq_lens_tensor.device,
418
dtype=attn_metadata.seq_lens_tensor.dtype)
419
causal_mask = False
@@ -584,6 +590,8 @@ class ROCmFlashAttentionImpl(AttentionImpl):
590
will match encoder sequence lengths, pass encoder sequence
591
attributes to kernel (encoder_seq_lens/encoder_seq_lens_tensor/
592
max_encoder_seq_len)
593
+ * ENCODER_ONLY: bidirectional attention with no KV caching;
594
+ use prefill sequence attributes
595
596
Args:
597
query: shape = [num_tokens, num_heads * head_size]
@@ -608,7 +616,11 @@ class ROCmFlashAttentionImpl(AttentionImpl):
616
else:
617
assert value is None
618
611
- if self.attn_type != AttentionType.ENCODER and kv_cache.numel() > 0:
619
+ # Only update KV cache for decoder self-attention
620
+ # and encoder-decoder cross-attention
621
+ if self.attn_type not in [
622
+ AttentionType.ENCODER, AttentionType.ENCODER_ONLY
623
+ ] and kv_cache.numel() > 0:
624
key_cache, value_cache = PagedAttention.split_kv_cache(
625
kv_cache, self.num_kv_heads, self.head_size)
626
@@ -632,6 +644,9 @@ class ROCmFlashAttentionImpl(AttentionImpl):
644
645
if self.attn_type != AttentionType.ENCODER:
646
num_prefill_tokens = attn_metadata.num_prefill_tokens
647
+ elif self.attn_type == AttentionType.ENCODER_ONLY:
648
+ # For encoder-only models, all tokens are processed in one go
649
+ num_prefill_tokens = query.shape[0]
650
else:
651
assert attn_metadata.num_encoder_tokens is not None
652
num_prefill_tokens = attn_metadata.num_encoder_tokens
@@ -642,8 +657,13 @@ class ROCmFlashAttentionImpl(AttentionImpl):
657
# QKV for prefill.
658
query = query[:num_prefill_tokens]
659
660
+ # For encoder-only and encoder models,
661
+ # we process all tokens at once
662
+ # For decoder and encoder-decoder,
663
+ # we may need to limit key/value to prefill tokens
664
if key is not None and value is not None \
646
- and self.attn_type != AttentionType.ENCODER_DECODER:
665
+ and self.attn_type not in [AttentionType.ENCODER_DECODER,
666
+ AttentionType.ENCODER_ONLY]:
667
key = key[:num_prefill_tokens]
668
value = value[:num_prefill_tokens]
669
@@ -678,7 +698,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
698
self.alibi_slopes,
699
query.dtype,
700
seq_lens,
681
- make_attn_mask=False) # type: ignore
701
+ make_attn_mask=causal_mask) # type: ignore
702
out, _ = self.attn_func(
703
query,
704
key,
@@ -703,7 +723,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
723
self.alibi_slopes,
724
query.dtype,
725
attn_metadata.seq_lens,
706
- make_attn_mask=True) # type: ignore
726
+ make_attn_mask=causal_mask) # type: ignore
727
query = query.movedim(0, query.dim() - 2)
728
key = key.movedim(0, key.dim() - 2)
729
value = value.movedim(0, value.dim() - 2)
@@ -729,7 +749,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
749
max_seqlen_q=prefill_meta.max_prefill_seq_len,
750
max_seqlen_k=key_max_seq_len,
751
softmax_scale=self.scale,
732
- causal=True,
752
+ causal=causal_mask,
753
window_size=self.sliding_window,
754
alibi_slopes=self.alibi_slopes,
755
softcap=self.logits_soft_cap,
@@ -742,25 +762,29 @@ class ROCmFlashAttentionImpl(AttentionImpl):
762
else:
763
output = out
764
else:
745
- # prefix-enabled attention
746
- output[:num_prefill_tokens] = PagedAttention.forward_prefix(
747
- query,
748
- key,
749
- value,
750
- self.kv_cache_dtype,
751
- key_cache,
752
- value_cache,
753
- prefill_meta.block_tables,
754
- prefill_meta.query_start_loc,
755
- prefill_meta.seq_lens_tensor,
756
- prefill_meta.max_query_len,
757
- self.alibi_slopes,
758
- self.sliding_window[0],
759
- layer._k_scale,
760
- layer._v_scale,
761
- )
762
-
763
- if decode_meta := attn_metadata.decode_metadata:
765
+ # prefix-enabled attention -
766
+ # not applicable for encoder-only models
767
+ if self.attn_type != AttentionType.ENCODER_ONLY:
768
+ output[:
769
+ num_prefill_tokens] = PagedAttention.forward_prefix(
770
+ query,
771
+ key,
772
+ value,
773
+ self.kv_cache_dtype,
774
+ key_cache,
775
+ value_cache,
776
+ prefill_meta.block_tables,
777
+ prefill_meta.query_start_loc,
778
+ prefill_meta.seq_lens_tensor,
779
+ prefill_meta.max_query_len,
780
+ self.alibi_slopes,
781
+ self.sliding_window[0],
782
+ layer._k_scale,
783
+ layer._v_scale,
784
+ )
785
+ # Skip decode phase for encoder-only models
786
+ if (decode_meta := attn_metadata.decode_metadata) and (
787
+ self.attn_type != AttentionType.ENCODER_ONLY):
788
# Decoding run.
789
# Whether to use rocm custom paged attention or not
790
num_seqs, num_heads, head_size = decode_query.shape
@@ -885,4 +909,4 @@ def _use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
909
and (qtype == torch.half or qtype == torch.bfloat16)
910
and (head_size == 64 or head_size == 128)
911
and (block_size == 16 or block_size == 32)
888
- and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768)
\ No newline at end of file
912
+ and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768)