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 8, 14:13 UTC
Add flagos in MiniCPM-o (#34126)
Signed-off-by: tc-mb <caitianchi@modelbest.cn> Signed-off-by: Vincent-Xiao <vincent.xiao.me@gmail.com> Co-authored-by: Vincent-Xiao <vincent.xiao.me@gmail.com>
tc-mb committed
Feb 10, 2026 at 18:51 UTC
e042d7e685daacfa9d4df92cc7d330060327a32b
1 file changed
+42
vllm/model_executor/models/minicpmo.py
+42
@@ -24,6 +24,7 @@
24
# limitations under the License.
25
"""Inference-only MiniCPM-O model compatible with HuggingFace weights."""
26
27
+import os
28
from collections.abc import Callable, Iterable, Mapping, Sequence
29
from typing import Annotated, Any, Literal, TypeAlias
30
@@ -75,6 +76,47 @@ from .utils import AutoWeightsLoader, cast_overflow_tensors, maybe_prefix
76
77
CPU_DEVICE = torch.device("cpu")
78
79
+if os.getenv("USE_FLAGOS") == "1":
80
+ import flag_gems
81
+
82
+ FLAG_GEMS_CONFIG = [
83
+ "sort",
84
+ "sort_stable",
85
+ "layer_norm",
86
+ "clamp_",
87
+ "cos",
88
+ "embedding",
89
+ "exp",
90
+ "exponential_",
91
+ "full",
92
+ "gather",
93
+ "gelu",
94
+ "index",
95
+ "le",
96
+ "lt",
97
+ "lt_scalar",
98
+ "masked_fill_",
99
+ "max",
100
+ "ones",
101
+ "pow_scalar",
102
+ "prod_dim",
103
+ "rand_like",
104
+ "reciprocal",
105
+ "repeat",
106
+ "scatter",
107
+ "scatter_",
108
+ "sin",
109
+ "sub",
110
+ "true_divide",
111
+ "true_divide_",
112
+ "uniform_",
113
+ "where_scalar_self",
114
+ "where_self_out",
115
+ "zeros",
116
+ "zeros_like",
117
+ ]
118
+ flag_gems.only_enable(record=False, include=FLAG_GEMS_CONFIG)
119
+
120
121
class MiniCPMOAudioFeatureInputs(TensorSchema):
122
"""