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 14, 05:30 UTC
[Bugfix] Fix out of bound index issue for Jina-embedding-v3 RoPE with cuda graph (#26687)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Oct 13, 2025 at 18:21 UTC
8e67b2557aae7204c697d7a5c61e00754da465be
2 files changed
+13
-7
tests/models/language/pooling_mteb_test/test_jina.py
-4
@@ -25,10 +25,6 @@ EMBEDDING_MODELS = [
25
mteb_score=0.824413164,
26
architecture="XLMRobertaModel",
27
is_matryoshka=True,
28
- # The default max length of the model is 8194, which will crash
29
- # CUDAGraph due to odd length for Gemm. We set it to 8192 to avoid
30
- # avoid this issue.
31
- max_model_len=8192,
28
dtype="float32",
29
)
30
]
vllm/model_executor/models/config.py
+13
-3
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
6
import vllm.envs as envs
7
from vllm.logger import init_logger
8
from vllm.model_executor.models import ModelRegistry
9
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
9
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv, round_up
10
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
11
12
if TYPE_CHECKING:
@@ -59,16 +59,26 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
59
class JinaRobertaModelConfig(VerifyAndUpdateConfig):
60
@staticmethod
61
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
62
- config = vllm_config.model_config.hf_config
62
+ model_config = vllm_config.model_config
63
+ config = model_config.hf_config
64
65
if config.position_embedding_type == "rotary":
66
assert config.__class__.__name__ == "XLMRobertaFlashConfig"
67
68
head_dim = config.hidden_size // config.num_attention_heads
69
+ max_position = config.max_position_embeddings
70
+ # Jina-embeddings-v3 has max_position_embeddings=8194, which will cause
71
+ # out-of-bound index issue at RoPE for long prompts with torch.compile,
72
+ # because it can't be divided by triton num_warps(default=4 or 8).
73
+ # To deal with this, we increase max_position to multiple of n_warps,
74
+ # so that triton kernel won't hit out-of-bound index in RoPE cache.
75
+ if not model_config.enforce_eager:
76
+ max_position = round_up(max_position, 8)
77
+
78
config.rotary_kwargs = {
79
"head_size": head_dim,
80
"rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
71
- "max_position": config.max_position_embeddings,
81
+ "max_position": max_position,
82
"base": getattr(config, "rope_theta", config.rotary_emb_base),
83
"rope_scaling": getattr(config, "rope_scaling", None),
84
}