Migrate embeddings inference service from HF Transformers+CK FlashAttention to vLLM+AITER
Closed, ResolvedPublic
Actions

Assigned To

Authored By

	kevinbazira
	Mar 4 2026, 8:09 AM

Project Tags

Referenced Files

	F72574039: image (12).png
	Mar 5 2026, 8:54 AM

Subscribers

Aklapper

kevinbazira

OKarakaya-WMF

Description

In T412338, we implemented the embeddings model-server using the HuggingFace Transformers inference backend with CK FlashAttention. As we work towards standardizing our LLM serving backend (T385173, T415627), we are migrating this service to use vLLM with AITER, which will improve throughput and memory efficiency (via PagedAttention).

Details

Other Assignee: OKarakaya-WMF

Related Changes in Gerrit:

Subject	Author	Repo	Branch	Lines +/-
ml-services: update embeddings-staging image to one that installs the hipsolver-dev headers	Kevin Bazira	operations/deployment-charts	master	+1 -1
embeddings: install hipsolver-dev package	Kevin Bazira	machinelearning/liftwing/inference-services	main	+1 -0
ml-services: update embeddings-staging image to one that installs the hipblaslt-dev headers	Kevin Bazira	operations/deployment-charts	master	+1 -1
embeddings: install the hipBLASLt development headers	Kevin Bazira	machinelearning/liftwing/inference-services	main	+1 -0
ml-services: add VLLM_ROCM_USE_AITER performance optimization env var to embeddings-staging isvc	Kevin Bazira	operations/deployment-charts	master	+6 -0
ml-services: add MAX_NUM_BATCHED_TOKENS performance optimization env var to embeddings-staging isvc	Kevin Bazira	operations/deployment-charts	master	+2 -0
ml-services: add MAX_MODEL_LEN performance optimization env var to embeddings-staging isvc	Kevin Bazira	operations/deployment-charts	master	+2 -0
ml-services: remove unused env vars in embeddings-staging isvc	Kevin Bazira	operations/deployment-charts	master	+3 -7
ml-services: add embeddings-staging isvc	Kevin Bazira	operations/deployment-charts	master	+50 -0
ml-services: rollback embeddings isvc image	Kevin Bazira	operations/deployment-charts	master	+6 -14
ml-services: revert embeddings isvc image to one that doesn't use AITER	Kevin Bazira	operations/deployment-charts	master	+2 -4
ml-services: bump llm limitranges to enable embeddings isvc deployment	Kevin Bazira	operations/deployment-charts	master	+21 -0
ml-services: increase memory in embeddings isvc to fix OOM issue caused AITER	Kevin Bazira	operations/deployment-charts	master	+6 -2
ml-services: update embeddings isvc to image that adds missing dev headers required by AITER	Kevin Bazira	operations/deployment-charts	master	+1 -1
embeddings: install ROCm libs to add missing dev headers required by AITER	Kevin Bazira	machinelearning/liftwing/inference-services	main	+6 -0
ml-services: update embeddings isvc to image that includes GPU device bitcode files	Kevin Bazira	operations/deployment-charts	master	+5 -1
embeddings: install ROCm device libs to fix runtime compilation failure	Kevin Bazira	machinelearning/liftwing/inference-services	main	+11 -0
ml-services: add performance optimization env vars to embeddings isvc	Kevin Bazira	operations/deployment-charts	master	+8 -6
ml-services: update embeddings isvc image	Kevin Bazira	operations/deployment-charts	master	+1 -1
embeddings: migrate inference backend from transformers+fa2 to vLLM	Kevin Bazira	machinelearning/liftwing/inference-services	main	+123 -131

Related Objects
Search...

Status	Assigned	Task
Resolved	OKarakaya-WMF	T412338 Q2 FY2025-26 Goal: Semantic Search - Embeddings Service for MVP
Resolved	kevinbazira	T418976 Migrate embeddings inference service from HF Transformers+CK FlashAttention to vLLM+AITER
Resolved	kevinbazira	T419650 Add ROCm build dependencies to wmf-debian-vllm image to support AITER kernel compilation

Event Timeline

There are a very large number of changes, so older changes are hidden. Show Older Changes

kevinbazira mentioned this in rMLISba253cb5f94d: embeddings: migrate inference backend from transformers+fa2 to vLLM.Mar 4 2026, 2:51 PM

Maintenance_bot removed a project: Patch-For-Review.Mar 4 2026, 3:33 PM

The embeddings model-server's inference backend has been migrated from transformers+fa2 to vLLM 0.14. I've tested the new model-server on ml-lab, and it was able to load the Qwen3-Embedding-0.6B model using vLLM on an MI210 GPU and generate embeddings as shown below:

1.Run the model-server

$ docker run --network=host -it \
--user root \
--device=/dev/kfd --device=/dev/dri \
--group-add=$(getent group video | cut -d: -f3) \
--group-add=$(getent group render | cut -d: -f3) \
--ipc=host \
--security-opt seccomp=unconfined \
--shm-size 16g \
-v /home/kevinbazira/test_embeddings/qwen3-embedding-model:/mnt/models \
docker-registry.wikimedia.org/wikimedia/machinelearning-liftwing-inference-services-embeddings:2026-03-04-145043-publish

It should start the inference-service as shown in:

P89798 Qwen3-Embedding-0.6B LLM loaded successfully in embeddings isvc that has been migrated from transformers+fa2 to vLLM

1	+ source common_settings.sh
2	+++ /usr/bin/python3 -c 'from python.resource_utils import get_cpu_count; print(get_cpu_count())'
3	++ CPU_COUNT=96
4	++ echo 'CPU count detected from get_cpu_count: 96'
5	CPU count detected from get_cpu_count: 96
6	++ export OMP_NUM_THREADS=96
7	++ OMP_NUM_THREADS=96
8	++ echo 'OMP_NUM_THREADS set to: 96'
9	OMP_NUM_THREADS set to: 96
10	+ MODEL_SERVER_PATH=model_server/model.py
11	+ exec /usr/bin/python3 model_server/model.py
12	/opt/amdgpu/share/libdrm/amdgpu.ids: No such file or directory
13	/opt/amdgpu/share/libdrm/amdgpu.ids: No such file or directory
14	[2026-03-04 15:21:35] INFO model.py:41: Loading vLLM model...
15	INFO 03-04 15:21:35 [utils.py:282] non-default args: {'trust_remote_code': True, 'dtype': 'float16', 'max_model_len': 8192, 'disable_log_stats': True, 'model': '/mnt/models/'}
16	The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
17	The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
18	[aiter] import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
19	[2026-03-04 15:21:43] INFO core.py:477: import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
20	INFO 03-04 15:21:43 [config.py:769] Found sentence-transformers modules configuration.
21	INFO 03-04 15:21:43 [config.py:796] Found pooling configuration.
22	INFO 03-04 15:21:43 [model.py:812] Resolved `--runner auto` to `--runner pooling`. Pass the value explicitly to silence this message.
23	INFO 03-04 15:21:43 [model.py:871] Resolved `--convert auto` to `--convert embed`. Pass the value explicitly to silence this message.
24	INFO 03-04 15:21:43 [model.py:541] Resolved architecture: Qwen3ForCausalLM
25	WARNING 03-04 15:21:43 [model.py:1885] Casting torch.bfloat16 to torch.float16.
26	INFO 03-04 15:21:43 [model.py:1561] Using max model len 8192
27	INFO 03-04 15:21:44 [scheduler.py:226] Chunked prefill is enabled with max_num_batched_tokens=8192.
28	INFO 03-04 15:21:44 [vllm.py:624] Asynchronous scheduling is enabled.
29	WARNING 03-04 15:21:44 [vllm.py:741] Pooling models do not support full cudagraphs. Overriding cudagraph_mode to PIECEWISE.
30	WARNING 03-04 15:21:44 [system_utils.py:140] We must use the `spawn` multiprocessing start method. Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. See https://docs.vllm.ai/en/latest/usage/troubleshooting.html#python-multiprocessing for more information. Reasons: CUDA is initialized
31	/opt/amdgpu/share/libdrm/amdgpu.ids: No such file or directory
32	/opt/amdgpu/share/libdrm/amdgpu.ids: No such file or directory
33	(EngineCore_DP0 pid=342) INFO 03-04 15:21:49 [core.py:96] Initializing a V1 LLM engine (v0.14.0rc2.dev374+g6c0064571) with config: model='/mnt/models/', speculative_config=None, tokenizer='/mnt/models/', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.float16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, disable_custom_all_reduce=True, quantization=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=cuda, structured_outputs_config=StructuredOutputsConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_parser='', reasoning_parser_plugin='', enable_in_reasoning=False), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, kv_cache_metrics=False, kv_cache_metrics_sample=0.01, cudagraph_metrics=False, enable_layerwise_nvtx_tracing=False, enable_mfu_metrics=False, enable_mm_processor_stats=False, enable_logging_iteration_details=False), seed=0, served_model_name=/mnt/models/, enable_prefix_caching=True, enable_chunked_prefill=True, pooler_config=PoolerConfig(pooling_type=None, seq_pooling_type='LAST', tok_pooling_type='ALL', normalize=None, dimensions=None, enable_chunked_processing=None, max_embed_len=None, softmax=None, activation=None, use_activation=True, logit_bias=None, step_tag_id=None, returned_token_ids=None), compilation_config={'level': None, 'mode': <CompilationMode.VLLM_COMPILE: 3>, 'debug_dump_path': None, 'cache_dir': '', 'compile_cache_save_format': 'binary', 'backend': 'inductor', 'custom_ops': ['none', '+sparse_attn_indexer'], 'splitting_ops': ['vllm::unified_attention', 'vllm::unified_attention_with_output', 'vllm::unified_mla_attention', 'vllm::unified_mla_attention_with_output', 'vllm::mamba_mixer2', 'vllm::mamba_mixer', 'vllm::short_conv', 'vllm::linear_attention', 'vllm::plamo2_mamba_mixer', 'vllm::gdn_attention_core', 'vllm::kda_attention', 'vllm::sparse_attn_indexer', 'vllm::rocm_aiter_sparse_attn_indexer'], 'compile_mm_encoder': False, 'compile_sizes': [], 'compile_ranges_split_points': [8192], 'inductor_compile_config': {'enable_auto_functionalized_v2': False, 'combo_kernels': True, 'benchmark_combo_kernel': True}, 'inductor_passes': {}, 'cudagraph_mode': <CUDAGraphMode.PIECEWISE: 1>, 'cudagraph_num_of_warmups': 1, 'cudagraph_capture_sizes': [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'cudagraph_copy_inputs': False, 'cudagraph_specialize_lora': True, 'use_inductor_graph_partition': False, 'pass_config': {'fuse_norm_quant': False, 'fuse_act_quant': False, 'fuse_attn_quant': False, 'eliminate_noops': True, 'enable_sp': False, 'fuse_gemm_comms': False, 'fuse_allreduce_rms': False}, 'max_cudagraph_capture_size': 512, 'dynamic_shapes_config': {'type': <DynamicShapesType.BACKED: 'backed'>, 'evaluate_guards': False, 'assume_32_bit_indexing': False}, 'local_cache_dir': None}
34	(EngineCore_DP0 pid=342) [aiter] import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
35	(EngineCore_DP0 pid=342) [2026-03-04 15:21:51] INFO core.py:477: import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
36	(EngineCore_DP0 pid=342) INFO 03-04 15:21:51 [parallel_state.py:1234] world_size=1 rank=0 local_rank=0 distributed_init_method=tcp://10.64.161.5:39109 backend=nccl
37	(EngineCore_DP0 pid=342) INFO 03-04 15:21:51 [parallel_state.py:1445] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, PCP rank 0, TP rank 0, EP rank N/A
38	(EngineCore_DP0 pid=342) INFO 03-04 15:21:52 [gpu_model_runner.py:4021] Starting to load model /mnt/models/...
39	(EngineCore_DP0 pid=342) INFO 03-04 15:21:53 [rocm.py:338] Using Triton Attention backend.
40	(EngineCore_DP0 pid=342) WARNING 03-04 15:21:53 [compilation.py:1048] Op 'sparse_attn_indexer' not present in model, enabling with '+sparse_attn_indexer' has no effect
41	Loading safetensors checkpoint shards: 0% Completed \| 0/1 [00:00<?, ?it/s]
42	(EngineCore_DP0 pid=342) INFO 03-04 15:21:53 [adapters.py:204] Mapping weights to Qwen3Model as they are relative to this model instead of Qwen3ForEmbedding.
43	Loading safetensors checkpoint shards: 100% Completed \| 1/1 [00:03<00:00, 3.97s/it]
44	Loading safetensors checkpoint shards: 100% Completed \| 1/1 [00:03<00:00, 3.97s/it]
45	(EngineCore_DP0 pid=342)
46	(EngineCore_DP0 pid=342) INFO 03-04 15:21:58 [default_loader.py:291] Loading weights took 4.77 seconds
47	(EngineCore_DP0 pid=342) INFO 03-04 15:21:58 [gpu_model_runner.py:4118] Model loading took 1.21 GiB memory and 5.445861 seconds
48	(EngineCore_DP0 pid=342) INFO 03-04 15:22:04 [backends.py:805] Using cache directory: /home/somebody/.cache/vllm/torch_compile_cache/73fb88de5c/rank_0_0/backbone for vLLM's torch.compile
49	(EngineCore_DP0 pid=342) INFO 03-04 15:22:04 [backends.py:865] Dynamo bytecode transform time: 5.19 s
50	(EngineCore_DP0 pid=342) INFO 03-04 15:22:14 [backends.py:302] Cache the graph of compile range (1, 8192) for later use
51	(EngineCore_DP0 pid=342) INFO 03-04 15:22:28 [backends.py:319] Compiling a graph for compile range (1, 8192) takes 21.70 s
52	(EngineCore_DP0 pid=342) INFO 03-04 15:22:28 [monitor.py:34] torch.compile takes 26.89 s in total
53	(EngineCore_DP0 pid=342) INFO 03-04 15:22:28 [decorators.py:576] saving AOT compiled function to /home/somebody/.cache/vllm/torch_aot_compile/b64e5f33a898c4b3c91b4e864a8e1d4429149f94df39e510dd24339e4b2979d6/rank_0_0/model
54	(EngineCore_DP0 pid=342) INFO 03-04 15:22:30 [decorators.py:580] saved AOT compiled function to /home/somebody/.cache/vllm/torch_aot_compile/b64e5f33a898c4b3c91b4e864a8e1d4429149f94df39e510dd24339e4b2979d6/rank_0_0/model
55	(EngineCore_DP0 pid=342) INFO 03-04 15:22:31 [gpu_worker.py:356] Available KV cache memory: 55.66 GiB
56	(EngineCore_DP0 pid=342) INFO 03-04 15:22:31 [kv_cache_utils.py:1307] GPU KV cache size: 521,120 tokens
57	(EngineCore_DP0 pid=342) INFO 03-04 15:22:31 [kv_cache_utils.py:1312] Maximum concurrency for 8,192 tokens per request: 63.61x
58	Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 100%\|████████████████████████████████████████████████████████████████████████\| 51/51 [00:01<00:00, 26.28it/s]
59	(EngineCore_DP0 pid=342) INFO 03-04 15:22:34 [gpu_model_runner.py:5051] Graph capturing finished in 3 secs, took 0.30 GiB
60	(EngineCore_DP0 pid=342) INFO 03-04 15:22:34 [core.py:272] init engine (profile, create kv cache, warmup model) took 35.43 seconds
61	(EngineCore_DP0 pid=342) INFO 03-04 15:22:34 [vllm.py:624] Asynchronous scheduling is enabled.
62	INFO 03-04 15:22:34 [llm.py:343] Supported tasks: ['token_embed', 'embed']
63	[2026-03-04 15:22:34] INFO model.py:56: vLLM model loaded successfully!
64	2026-03-04 15:22:34.987 1 kserve INFO [model_server.py:register_model():406] Registering model: qwen3-embedding
65	2026-03-04 15:22:34.987 1 kserve INFO [model_server.py:setup_event_loop():286] Setting max asyncio worker threads as 32
66	2026-03-04 15:22:35.069 1 kserve INFO [server.py:_register_endpoints():111] OpenAI endpoints not registered
67	2026-03-04 15:22:35.070 1 kserve INFO [server.py:_register_endpoints():119] Time series endpoints not registered
68	2026-03-04 15:22:35.070 1 kserve INFO [server.py:start():170] Starting uvicorn with 1 workers
69	2026-03-04 15:22:35.087 1 uvicorn.error INFO: Started server process [1]
70	2026-03-04 15:22:35.087 1 uvicorn.error INFO: Waiting for application startup.
71	2026-03-04 15:22:35.093 1 kserve INFO [server.py:start():70] Starting gRPC server with 4 workers
72	2026-03-04 15:22:35.093 1 kserve INFO [server.py:start():71] Starting gRPC server on [::]:8081
73	2026-03-04 15:22:35.093 1 uvicorn.error INFO: Application startup complete.
74	2026-03-04 15:22:35.093 1 uvicorn.error INFO: Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

2.Query the inference service

$ time curl -s localhost:8080/v1/models/qwen3-embedding:predict -X POST -d '{"input": ["text1", "text2"]}' -i -H "Content-type: application/json"

HTTP/1.1 200 OK
date: Wed, 04 Mar 2026 15:26:03 GMT
server: uvicorn
content-length: 43502
content-type: application/json

{
    "object":"list",
    "data":[
        {
            "object":"embedding",
            "embedding":[-0.036402568221092224, ...],
            "index":0
        },
        {
            "object":"embedding",
            "embedding":[-0.025051085278391838, ...],
            "index":1
        }
    ],
    "model":"qwen3-embedding"
}

real	0m0.789s
user	0m0.005s
sys	0m0.009s

Change #1248310 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: update embeddings isvc image

https://gerrit.wikimedia.org/r/1248310

gerritbot added a project: Patch-For-Review.Mar 5 2026, 6:43 AM

Change #1248310 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: update embeddings isvc image

https://gerrit.wikimedia.org/r/1248310

The new embeddings model-server that uses vLLM as the inference backend instead of transformers+fa2 has been deployed in LiftWing prod on an MI300x GPU:

$ time curl "https://inference.svc.eqiad.wmnet:30443/v1/models/qwen3-embedding:predict" -X POST -d '{"input": ["text1", "text2"]}' -H "Host: embeddings.llm.wikimedia.org" -H "Content-Type: application/json" --http1.1

{
    "object":"list",
    "data":[
        {
            "object":"embedding",
            "embedding":[-0.03630506619811058, ...],
            "index":0
        },
        {
            "object":"embedding",
            "embedding":[-0.02502359449863434, ...],
            "index":1
        }
    ],
    "model":"Qwen3-Embedding-0.6B"
}

real	0m0.482s
user	0m0.016s
sys	0m0.003s

Maintenance_bot removed a project: Patch-For-Review.Mar 5 2026, 8:31 AM

transformer test1 (light):

(venv) ozge@stat1010:~/repos/wiki/gerrit/inference-services/test/locust$ MODEL=embeddings locust
Min length: 50, Max length: 350
       question_length
count      3952.000000
mean         84.550354
std          37.088196
min          50.000000
25%          64.000000
50%          77.000000
75%          93.000000
max         348.000000
Prompt length:  108
[2026-01-08 12:50:39,141] stat1010/INFO/locust.main: Run time limit set to 120 seconds
[2026-01-08 12:50:39,142] stat1010/INFO/locust.main: Starting Locust 2.31.5
[2026-01-08 12:50:39,142] stat1010/INFO/locust.runners: Ramping to 2 users at a rate of 10.00 per second
[2026-01-08 12:50:39,143] stat1010/INFO/locust.runners: All users spawned: {"Embeddings": 2} (2 total users)
[2026-01-08 12:52:38,649] stat1010/INFO/locust.main: --run-time limit reached, shutting down
Load test results are within the threshold
[2026-01-08 12:52:38,727] stat1010/INFO/locust.main: Shutting down (exit code 0)
Type     Name                                                                          # reqs      # fails |    Avg     Min     Max    Med |   req/s  failures/s
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
POST     /v1/models/qwen3-embedding:predict                                              2586     0(0.00%) |     41      28     292     34 |   21.64        0.00
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
         Aggregated                                                                      2586     0(0.00%) |     41      28     292     34 |   21.64        0.00

Response time percentiles (approximated)
Type     Name                                                                                  50%    66%    75%    80%    90%    95%    98%    99%  99.9% 99.99%   100% # reqs
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
POST     /v1/models/qwen3-embedding:predict                                                     34     38     45     48     56     61     81    260    280    290    290   2586
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
         Aggregated

vllm test1 (light):

(venv) ozge@stat1010:~/repos/wiki/gerrit/inference-services/test/locust$ MODEL=embeddings locust --host https://inference.svc.eqiad.wmnet:30443
Min length: 50, Max length: 350
       question_length
count      3952.000000
mean         84.550354
std          37.088196
min          50.000000
25%          64.000000
50%          77.000000
75%          93.000000
max         348.000000
Prompt length:  108
[2026-03-05 08:41:19,403] stat1010/INFO/locust.main: Run time limit set to 120 seconds
[2026-03-05 08:41:19,403] stat1010/INFO/locust.main: Starting Locust 2.31.5
[2026-03-05 08:41:19,403] stat1010/INFO/locust.runners: Ramping to 2 users at a rate of 10.00 per second
[2026-03-05 08:41:19,404] stat1010/INFO/locust.runners: All users spawned: {"Embeddings": 2} (2 total users)
[2026-03-05 08:43:18,886] stat1010/INFO/locust.main: --run-time limit reached, shutting down
Load test results are within the threshold
[2026-03-05 08:43:18,973] stat1010/INFO/locust.main: Shutting down (exit code 0)
Type     Name                                                                          # reqs      # fails |    Avg     Min     Max    Med |   req/s  failures/s
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
POST     /v1/models/qwen3-embedding:predict                                              3454     0(0.00%) |     18      12     757     15 |   28.89        0.00
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
         Aggregated                                                                      3454     0(0.00%) |     18      12     757     15 |   28.89        0.00

Response time percentiles (approximated)
Type     Name                                                                                  50%    66%    75%    80%    90%    95%    98%    99%  99.9% 99.99%   100% # reqs
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
POST     /v1/models/qwen3-embedding:predict                                                     15     16     17     17     21     24     26     31    240    760    760   3454
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
         Aggregated                                                                             15     16     17     17     21     24     26     31    240    760    760   3454

transformer test2 (heavy):

vllm test2 (heavy):

(venv) ozge@stat1010:~/repos/wiki/gerrit/inference-services/test/locust$ MODEL=embeddings locust --host https://inference.svc.eqiad.wmnet:30443 --users 10 --run-time 4m
Min length: 10, Max length: 350
       question_length
count      4610.000000
mean         78.490456
std          37.470874
min          20.000000
25%          58.000000
50%          73.000000
75%          90.000000
max         348.000000
Prompt length:  108
[2026-03-05 08:44:37,699] stat1010/INFO/locust.main: Run time limit set to 240 seconds
[2026-03-05 08:44:37,699] stat1010/INFO/locust.main: Starting Locust 2.31.5
[2026-03-05 08:44:37,700] stat1010/INFO/locust.runners: Ramping to 10 users at a rate of 10.00 per second
[2026-03-05 08:44:37,700] stat1010/INFO/locust.runners: All users spawned: {"Embeddings": 10} (10 total users)
[2026-03-05 08:48:37,193] stat1010/INFO/locust.main: --run-time limit reached, shutting down
Load test results are within the threshold
[2026-03-05 08:48:37,292] stat1010/INFO/locust.main: Shutting down (exit code 1)
Type     Name                                                                          # reqs      # fails |    Avg     Min     Max    Med |   req/s  failures/s
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
POST     /v1/models/qwen3-embedding:predict                                             22610     1(0.00%) |     54      12     573     52 |   94.38        0.00
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
         Aggregated                                                                     22610     1(0.00%) |     54      12     573     52 |   94.38        0.00

Response time percentiles (approximated)
Type     Name                                                                                  50%    66%    75%    80%    90%    95%    98%    99%  99.9% 99.99%   100% # reqs
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
POST     /v1/models/qwen3-embedding:predict                                                     52     58     62     64     71     77     88    250    320    340    570  22610
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
         Aggregated                                                                             52     58     62     64     71     77     88    250    320    340    570  22610

Error report
# occurrences      Error
------------------|---------------------------------------------------------------------------------------------------------------------------------------------
1                  POST /v1/models/qwen3-embedding:predict: BadStatusCode('https://inference.svc.eqiad.wmnet:30443/v1/models/qwen3-embedding:predict', code=502)
------------------|---------------------------------------------------------------------------------------------------------------------------------------------

Overall, number of requests per second increased by 3X, and median latency decreased by 4.5X

Change #1248416 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: add performance optimization env vars to embeddings isvc

https://gerrit.wikimedia.org/r/1248416

gerritbot added a project: Patch-For-Review.Mar 5 2026, 10:08 AM

Change #1248416 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: add performance optimization env vars to embeddings isvc

https://gerrit.wikimedia.org/r/1248416

Maintenance_bot removed a project: Patch-For-Review.Mar 5 2026, 11:30 AM

Change #1248457 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[machinelearning/liftwing/inference-services@main] embeddings: install ROCm device libs to fix runtime compilation failure

https://gerrit.wikimedia.org/r/1248457

Change #1248457 merged by jenkins-bot:

[machinelearning/liftwing/inference-services@main] embeddings: install ROCm device libs to fix runtime compilation failure

https://gerrit.wikimedia.org/r/1248457

kevinbazira mentioned this in rMLIS6f76f1391818: embeddings: install ROCm device libs to fix runtime compilation failure.Mar 5 2026, 12:29 PM

Maintenance_bot removed a project: Patch-For-Review.Mar 5 2026, 12:31 PM

Change #1248462 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: update embeddings isvc to image that includes GPU device bitcode files

https://gerrit.wikimedia.org/r/1248462

gerritbot added a project: Patch-For-Review.Mar 5 2026, 12:51 PM

Change #1248462 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: update embeddings isvc to image that includes GPU device bitcode files

https://gerrit.wikimedia.org/r/1248462

Maintenance_bot removed a project: Patch-For-Review.Mar 5 2026, 1:31 PM

Change #1248475 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[machinelearning/liftwing/inference-services@main] embeddings: install ROCm libs to add missing dev headers required by AITER

https://gerrit.wikimedia.org/r/1248475

gerritbot added a project: Patch-For-Review.Mar 5 2026, 1:34 PM

Change #1248475 merged by jenkins-bot:

[machinelearning/liftwing/inference-services@main] embeddings: install ROCm libs to add missing dev headers required by AITER

https://gerrit.wikimedia.org/r/1248475

kevinbazira mentioned this in rMLISa13f23b783e9: embeddings: install ROCm libs to add missing dev headers required by AITER.Mar 5 2026, 1:39 PM

Change #1248480 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: update embeddings isvc to image that adds missing dev headers required by AITER

https://gerrit.wikimedia.org/r/1248480

Change #1248480 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: update embeddings isvc to image that adds missing dev headers required by AITER

https://gerrit.wikimedia.org/r/1248480

Maintenance_bot removed a project: Patch-For-Review.Mar 5 2026, 3:31 PM

Change #1248522 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: increase memory in embeddings isvc to fix OOM issue caused AITER

https://gerrit.wikimedia.org/r/1248522

gerritbot added a project: Patch-For-Review.Mar 5 2026, 4:05 PM

Change #1248522 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: increase memory in embeddings isvc to fix OOM issue caused AITER

https://gerrit.wikimedia.org/r/1248522

Maintenance_bot removed a project: Patch-For-Review.Mar 5 2026, 4:31 PM

Change #1248531 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: bump llm limitranges to enable embeddings isvc deployment

https://gerrit.wikimedia.org/r/1248531

gerritbot added a project: Patch-For-Review.Mar 5 2026, 4:43 PM

EBernhardson mentioned this in T419174: qwen3-embedding:predict returning 503 to all requests.Mar 5 2026, 10:05 PM

Change #1248531 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: bump llm limitranges to enable embeddings isvc deployment

https://gerrit.wikimedia.org/r/1248531

Maintenance_bot removed a project: Patch-For-Review.Mar 6 2026, 10:31 AM

Change #1248855 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: revert embeddings isvc image to one that doesn't use AITER

https://gerrit.wikimedia.org/r/1248855

gerritbot added a project: Patch-For-Review.Mar 6 2026, 3:09 PM

Change #1248855 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: revert embeddings isvc image to one that doesn't use AITER

https://gerrit.wikimedia.org/r/1248855

Change #1248859 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: rollback embeddings isvc image

https://gerrit.wikimedia.org/r/1248859

Change #1248859 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: rollback embeddings isvc image

https://gerrit.wikimedia.org/r/1248859

Maintenance_bot removed a project: Patch-For-Review.Mar 6 2026, 3:30 PM

Change #1249081 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: add embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249081

gerritbot added a project: Patch-For-Review.Mar 9 2026, 5:46 AM

Change #1249081 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: add embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249081

Maintenance_bot removed a project: Patch-For-Review.Mar 9 2026, 8:30 AM

Change #1249204 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: remove unused env vars in embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249204

gerritbot added a project: Patch-For-Review.Mar 9 2026, 8:50 AM

Change #1249204 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: remove unused env vars in embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249204

Change #1249209 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: add MAX_MODEL_LEN performance optimization env var to embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249209

Change #1249209 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: add MAX_MODEL_LEN performance optimization env var to embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249209

Maintenance_bot removed a project: Patch-For-Review.Mar 9 2026, 9:30 AM

Change #1249213 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: add MAX_NUM_BATCHED_TOKENS performance optimization env var to embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249213

gerritbot added a project: Patch-For-Review.Mar 9 2026, 9:44 AM

Change #1249213 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: add MAX_NUM_BATCHED_TOKENS performance optimization env var to embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249213

Maintenance_bot removed a project: Patch-For-Review.Mar 9 2026, 10:31 AM

Change #1249255 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: add VLLM_ROCM_USE_AITER performance optimization env var to embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249255

gerritbot added a project: Patch-For-Review.Mar 9 2026, 11:45 AM

Change #1249255 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: add VLLM_ROCM_USE_AITER performance optimization env var to embeddings-staging isvc

https://gerrit.wikimedia.org/r/1249255

Maintenance_bot removed a project: Patch-For-Review.Mar 9 2026, 1:33 PM

Change #1249326 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[machinelearning/liftwing/inference-services@main] embeddings: install the hipBLASLt development headers

https://gerrit.wikimedia.org/r/1249326

gerritbot added a project: Patch-For-Review.Mar 9 2026, 2:43 PM

Change #1249326 merged by jenkins-bot:

[machinelearning/liftwing/inference-services@main] embeddings: install the hipBLASLt development headers

https://gerrit.wikimedia.org/r/1249326

kevinbazira mentioned this in rMLISab73ca00785c: embeddings: install the hipBLASLt development headers.Mar 9 2026, 2:54 PM

Change #1249333 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: update embeddings-staging image to one that installs the hipblaslt-dev headers

https://gerrit.wikimedia.org/r/1249333

Change #1249333 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: update embeddings-staging image to one that installs the hipblaslt-dev headers

https://gerrit.wikimedia.org/r/1249333

Maintenance_bot removed a project: Patch-For-Review.Mar 9 2026, 3:31 PM

Change #1249338 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[machinelearning/liftwing/inference-services@main] embeddings: install hipsolver-dev package

https://gerrit.wikimedia.org/r/1249338

Change #1249338 merged by jenkins-bot:

[machinelearning/liftwing/inference-services@main] embeddings: install hipsolver-dev package

https://gerrit.wikimedia.org/r/1249338

kevinbazira mentioned this in rMLIS6af5dc133b2d: embeddings: install hipsolver-dev package.Mar 9 2026, 3:54 PM

Change #1249345 had a related patch set uploaded (by Kevin Bazira; author: Kevin Bazira):

[operations/deployment-charts@master] ml-services: update embeddings-staging image to one that installs the hipsolver-dev headers

https://gerrit.wikimedia.org/r/1249345

Change #1249345 merged by jenkins-bot:

[operations/deployment-charts@master] ml-services: update embeddings-staging image to one that installs the hipsolver-dev headers

https://gerrit.wikimedia.org/r/1249345

Maintenance_bot removed a project: Patch-For-Review.Mar 9 2026, 4:30 PM

vLLM uses Triton FlashAttention as the default attention backend. AMD released the AI Tensor Engine for ROCm (AITER), and vLLM now recommends using AITER as it's a more performant backend for MI300x GPUs. Based on this recommendation, we enabled AITER in the embeddings model-server.

On top of AITER, below are all the optimizations we have implemented:

VLLM_ROCM_USE_AITER: Enabled ROCm's AI Tensor Engine for MI300x specific matrix core optimizations.
VLLM_USE_TRITON_FLASH_ATTN: Disabled Triton for Attention in favor of AITER.
MAX_JOBS: Restricted the ninja build system to a single compilation thread. JIT-compiling the AITER kernels is memory-intensive, and multiple concurrent jobs can exceed the k8s pod's 16Gi RAM limit and trigger an OOM kill during the initial build.
MAX_MODEL_LEN: Matched Qwen3 sequence length limit to enable the search team to index entire articles without losing semantic information at the end of the document. We can always cap it at 8192, but any text beyond that point gets truncated (ignored) by the embedding model.
MAX_NUM_BATCHED_TOKENS: Matched the MAX_MODEL_LEN to ensure the engine can process at least one full-length article in a single pass, or efficiently pack hundreds of shorter search queries together. Leaving it at a lower default would bottleneck throughput and prevent us from utilizing the full 32K context in one go.
TRUST_REMOTE_CODE: Disabled in production since we store this model in Swift and prod doesn't need to access HuggingFace.

The embeddings-staging isvc starts in the pod on eqiad as shown below:

P89830 Qwen3-Embedding-0.6B LLM loaded successfully in embeddings isvc that uses vLLM+AITER on MI300x GPU

1	$ kubectl logs embeddings-staging-predictor-00007-deployment-f7b657987-zszsb
2	+ source common_settings.sh
3	+++ /usr/bin/python3 -c 'from python.resource_utils import get_cpu_count; print(get_cpu_count())'
4	++ CPU_COUNT=6
5	++ echo 'CPU count detected from get_cpu_count: 6'
6	++ export OMP_NUM_THREADS=6
7	++ OMP_NUM_THREADS=6
8	CPU count detected from get_cpu_count: 6
9	OMP_NUM_THREADS set to: 6
10	++ echo 'OMP_NUM_THREADS set to: 6'
11	+ MODEL_SERVER_PATH=model_server/model.py
12	+ exec /usr/bin/python3 model_server/model.py
13	/opt/amdgpu/share/libdrm/amdgpu.ids: No such file or directory
14	[2026-03-09 16:27:35] INFO model.py:43: Loading vLLM model...
15	INFO 03-09 16:27:35 [utils.py:282] non-default args: {'dtype': 'float16', 'max_model_len': 32768, 'max_num_batched_tokens': 32768, 'disable_log_stats': True, 'model': '/mnt/models/'}
16	[aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
17	[2026-03-09 16:27:47] WARNING core.py:454: WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
18	[aiter] import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
19	[2026-03-09 16:27:47] INFO core.py:477: import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
20	INFO 03-09 16:27:47 [config.py:769] Found sentence-transformers modules configuration.
21	INFO 03-09 16:27:47 [config.py:796] Found pooling configuration.
22	INFO 03-09 16:27:47 [model.py:812] Resolved `--runner auto` to `--runner pooling`. Pass the value explicitly to silence this message.
23	INFO 03-09 16:27:47 [model.py:871] Resolved `--convert auto` to `--convert embed`. Pass the value explicitly to silence this message.
24	INFO 03-09 16:27:47 [model.py:541] Resolved architecture: Qwen3ForCausalLM
25	WARNING 03-09 16:27:47 [model.py:1885] Casting torch.bfloat16 to torch.float16.
26	INFO 03-09 16:27:47 [model.py:1561] Using max model len 32768
27	INFO 03-09 16:27:48 [scheduler.py:226] Chunked prefill is enabled with max_num_batched_tokens=32768.
28	INFO 03-09 16:27:48 [vllm.py:624] Asynchronous scheduling is enabled.
29	WARNING 03-09 16:27:48 [vllm.py:741] Pooling models do not support full cudagraphs. Overriding cudagraph_mode to PIECEWISE.
30	WARNING 03-09 16:27:49 [system_utils.py:140] We must use the `spawn` multiprocessing start method. Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. See https://docs.vllm.ai/en/latest/usage/troubleshooting.html#python-multiprocessing for more information. Reasons: CUDA is initialized
31	/opt/amdgpu/share/libdrm/amdgpu.ids: No such file or directory
32	(EngineCore_DP0 pid=141) INFO 03-09 16:27:57 [core.py:96] Initializing a V1 LLM engine (v0.14.0rc2.dev374+g6c0064571) with config: model='/mnt/models/', speculative_config=None, tokenizer='/mnt/models/', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=32768, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=cuda, structured_outputs_config=StructuredOutputsConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_parser='', reasoning_parser_plugin='', enable_in_reasoning=False), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, kv_cache_metrics=False, kv_cache_metrics_sample=0.01, cudagraph_metrics=False, enable_layerwise_nvtx_tracing=False, enable_mfu_metrics=False, enable_mm_processor_stats=False, enable_logging_iteration_details=False), seed=0, served_model_name=/mnt/models/, enable_prefix_caching=True, enable_chunked_prefill=True, pooler_config=PoolerConfig(pooling_type=None, seq_pooling_type='LAST', tok_pooling_type='ALL', normalize=None, dimensions=None, enable_chunked_processing=None, max_embed_len=None, softmax=None, activation=None, use_activation=True, logit_bias=None, step_tag_id=None, returned_token_ids=None), compilation_config={'level': None, 'mode': <CompilationMode.VLLM_COMPILE: 3>, 'debug_dump_path': None, 'cache_dir': '', 'compile_cache_save_format': 'binary', 'backend': 'inductor', 'custom_ops': ['none', '+rms_norm', '+quant_fp8', '+grouped_topk', '+sparse_attn_indexer'], 'splitting_ops': ['vllm::unified_attention', 'vllm::unified_attention_with_output', 'vllm::unified_mla_attention', 'vllm::unified_mla_attention_with_output', 'vllm::mamba_mixer2', 'vllm::mamba_mixer', 'vllm::short_conv', 'vllm::linear_attention', 'vllm::plamo2_mamba_mixer', 'vllm::gdn_attention_core', 'vllm::kda_attention', 'vllm::sparse_attn_indexer', 'vllm::rocm_aiter_sparse_attn_indexer'], 'compile_mm_encoder': False, 'compile_sizes': [], 'compile_ranges_split_points': [32768], 'inductor_compile_config': {'enable_auto_functionalized_v2': False, 'combo_kernels': True, 'benchmark_combo_kernel': True}, 'inductor_passes': {}, 'cudagraph_mode': <CUDAGraphMode.PIECEWISE: 1>, 'cudagraph_num_of_warmups': 1, 'cudagraph_capture_sizes': [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'cudagraph_copy_inputs': False, 'cudagraph_specialize_lora': True, 'use_inductor_graph_partition': False, 'pass_config': {'fuse_norm_quant': False, 'fuse_act_quant': False, 'fuse_attn_quant': False, 'eliminate_noops': True, 'enable_sp': False, 'fuse_gemm_comms': False, 'fuse_allreduce_rms': False}, 'max_cudagraph_capture_size': 512, 'dynamic_shapes_config': {'type': <DynamicShapesType.BACKED: 'backed'>, 'evaluate_guards': False, 'assume_32_bit_indexing': False}, 'local_cache_dir': None}
33	(EngineCore_DP0 pid=141) [aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
34	(EngineCore_DP0 pid=141) [2026-03-09 16:27:58] WARNING core.py:454: WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
35	(EngineCore_DP0 pid=141) [aiter] import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
36	(EngineCore_DP0 pid=141) [2026-03-09 16:27:58] INFO core.py:477: import [module_aiter_enum] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_aiter_enum.so
37	(EngineCore_DP0 pid=141) INFO 03-09 16:27:59 [parallel_state.py:1234] world_size=1 rank=0 local_rank=0 distributed_init_method=tcp://10.67.19.59:42685 backend=nccl
38	(EngineCore_DP0 pid=141) INFO 03-09 16:27:59 [parallel_state.py:1445] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, PCP rank 0, TP rank 0, EP rank N/A
39	(EngineCore_DP0 pid=141) INFO 03-09 16:27:59 [topk_topp_sampler.py:77] Using aiter sampler on ROCm (lazy import, sampling-only).
40	(EngineCore_DP0 pid=141) INFO 03-09 16:27:59 [gpu_model_runner.py:4021] Starting to load model /mnt/models/...
41	(EngineCore_DP0 pid=141) INFO 03-09 16:28:00 [rocm.py:313] Using Aiter Flash Attention backend.
42	(EngineCore_DP0 pid=141) WARNING 03-09 16:28:00 [compilation.py:1048] Op 'quant_fp8' not present in model, enabling with '+quant_fp8' has no effect
43	(EngineCore_DP0 pid=141) WARNING 03-09 16:28:00 [compilation.py:1048] Op 'grouped_topk' not present in model, enabling with '+grouped_topk' has no effect
44	(EngineCore_DP0 pid=141) WARNING 03-09 16:28:00 [compilation.py:1048] Op 'sparse_attn_indexer' not present in model, enabling with '+sparse_attn_indexer' has no effect
45	Loading safetensors checkpoint shards: 0% Completed \| 0/1 [00:00<?, ?it/s]
46	(EngineCore_DP0 pid=141) INFO 03-09 16:28:00 [adapters.py:204] Mapping weights to Qwen3Model as they are relative to this model instead of Qwen3ForEmbedding.
47	Loading safetensors checkpoint shards: 100% Completed \| 1/1 [00:00<00:00, 136.68it/s]
48	(EngineCore_DP0 pid=141)
49	(EngineCore_DP0 pid=141) INFO 03-09 16:28:00 [default_loader.py:291] Loading weights took 0.15 seconds
50	(EngineCore_DP0 pid=141) INFO 03-09 16:28:01 [gpu_model_runner.py:4118] Model loading took 1.2 GiB memory and 0.644127 seconds
51	(EngineCore_DP0 pid=141) INFO 03-09 16:28:05 [backends.py:805] Using cache directory: /home/somebody/.cache/vllm/torch_compile_cache/2e85a3be17/rank_0_0/backbone for vLLM's torch.compile
52	(EngineCore_DP0 pid=141) INFO 03-09 16:28:05 [backends.py:865] Dynamo bytecode transform time: 3.71 s
53	(EngineCore_DP0 pid=141) INFO 03-09 16:28:11 [backends.py:302] Cache the graph of compile range (1, 32768) for later use
54	(EngineCore_DP0 pid=141) [aiter] import [module_rmsnorm] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_rmsnorm.so
55	(EngineCore_DP0 pid=141) [2026-03-09 16:28:12] INFO core.py:477: import [module_rmsnorm] under /srv/venv/lib/python3.11/site-packages/aiter/jit/module_rmsnorm.so
56	(EngineCore_DP0 pid=141) INFO 03-09 16:28:28 [backends.py:319] Compiling a graph for compile range (1, 32768) takes 21.72 s
57	(EngineCore_DP0 pid=141) INFO 03-09 16:28:28 [monitor.py:34] torch.compile takes 25.43 s in total
58	(EngineCore_DP0 pid=141) INFO 03-09 16:28:28 [decorators.py:576] saving AOT compiled function to /home/somebody/.cache/vllm/torch_aot_compile/35af5ebdfedbcf3c035a62ff4c8dcadaca6ad3d29fd79ab0b1b3dfc390b05c81/rank_0_0/model
59	(EngineCore_DP0 pid=141) INFO 03-09 16:28:30 [decorators.py:580] saved AOT compiled function to /home/somebody/.cache/vllm/torch_aot_compile/35af5ebdfedbcf3c035a62ff4c8dcadaca6ad3d29fd79ab0b1b3dfc390b05c81/rank_0_0/model
60	(EngineCore_DP0 pid=141) INFO 03-09 16:28:31 [gpu_worker.py:356] Available KV cache memory: 19.03 GiB
61	(EngineCore_DP0 pid=141) INFO 03-09 16:28:31 [kv_cache_utils.py:1307] GPU KV cache size: 178,112 tokens
62	(EngineCore_DP0 pid=141) INFO 03-09 16:28:31 [kv_cache_utils.py:1312] Maximum concurrency for 32,768 tokens per request: 5.44x
63	Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 100%\|██████████\| 51/51 [00:01<00:00, 26.61it/s]
64	(EngineCore_DP0 pid=141) INFO 03-09 16:28:33 [gpu_model_runner.py:5051] Graph capturing finished in 3 secs, took 0.30 GiB
65	(EngineCore_DP0 pid=141) INFO 03-09 16:28:34 [core.py:272] init engine (profile, create kv cache, warmup model) took 32.74 seconds
66	(EngineCore_DP0 pid=141) INFO 03-09 16:28:34 [vllm.py:624] Asynchronous scheduling is enabled.
67	INFO 03-09 16:28:34 [llm.py:343] Supported tasks: ['embed', 'token_embed']
68	[2026-03-09 16:28:34] INFO model.py:59: vLLM model loaded successfully!
69	2026-03-09 16:28:34.820 1 kserve INFO [model_server.py:register_model():406] Registering model: qwen3-embedding-staging
70	2026-03-09 16:28:34.821 1 kserve INFO [model_server.py:setup_event_loop():286] Setting max asyncio worker threads as 32
71	2026-03-09 16:28:34.849 1 kserve INFO [server.py:_register_endpoints():111] OpenAI endpoints not registered
72	2026-03-09 16:28:34.849 1 kserve INFO [server.py:_register_endpoints():119] Time series endpoints not registered
73	2026-03-09 16:28:34.849 1 kserve INFO [server.py:start():170] Starting uvicorn with 1 workers
74	2026-03-09 16:28:34.859 1 uvicorn.error INFO: Started server process [1]
75	2026-03-09 16:28:34.859 1 uvicorn.error INFO: Waiting for application startup.
76	2026-03-09 16:28:34.862 1 kserve INFO [server.py:start():70] Starting gRPC server with 4 workers
77	2026-03-09 16:28:34.862 1 kserve INFO [server.py:start():71] Starting gRPC server on [::]:8081
78	2026-03-09 16:28:34.862 1 uvicorn.error INFO: Application startup complete.
79	2026-03-09 16:28:34.862 1 uvicorn.error INFO: Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

Querying the embeddings-staging isvc returns:

$ time curl "https://inference.svc.eqiad.wmnet:30443/v1/models/qwen3-embedding-staging:predict" -X POST -d '{"input": ["text1", "text2"]}' -H "Host: embeddings-staging.llm.wikimedia.org" -H "Content-Type: application/json" --http1.1

{
    "object":"list",
    "data":[
        {
            "object":"embedding",
            "embedding":[-0.03621598333120346, ...],
            "index":0
        },
        {
            "object":"embedding","embedding":[-0.02509966306388378, ...],
            "index":1
        }
    ],
    "model":"Qwen3-Embedding-0.6B"
}

real	0m0.200s
user	0m0.017s
sys	0m0.004s

NOTE: The very first inference request triggers a JIT compilation of the AITER kernels. This results in a one-time startup delay as the system builds optimized GPU kernels for the MI300x (gfx942) hardware; however, once these kernels are compiled and cached, subsequent requests are processed at peak performance with significantly lower latency.

kevinbazira mentioned this in P89805 enbeddings isvc throws runtime compilation (JIT) failure when VLLM_ROCM_USE_AITER=1.Mar 10 2026, 7:13 AM

kevinbazira mentioned this in T385173: Use rocm/vllm image on Lift Wing.Mar 10 2026, 8:25 AM

kevinbazira renamed this task from Migrate embeddings inference service from Transformers+FA2 to vLLM to Migrate embeddings inference service from HF Transformers+CK FlashAttention to vLLM+AITER.Mar 11 2026, 5:11 AM

kevinbazira updated the task description. (Show Details)

kevinbazira created subtask T419650: Add ROCm build dependencies to wmf-debian-vllm image to support AITER kernel compilation.Mar 11 2026, 5:53 AM

kevinbazira mentioned this in T419650: Add ROCm build dependencies to wmf-debian-vllm image to support AITER kernel compilation.

achou moved this task from Unsorted to In Progress on the Machine-Learning-Team board.Mar 11 2026, 11:38 AM

after the aiter changes.
reference to compare is vllm test2 (heavy):

(venv) ozge@stat1010:~/repos/wiki/gerrit/inference-services/test/locust$ MODEL=embeddings locust --host https://inference.svc.eqiad.wmnet:30443 --users 10 --run-time 4m
Min length: 10, Max length: 350
       question_length
count      4610.000000
mean         78.490456
std          37.470874
min          20.000000
25%          58.000000
50%          73.000000
75%          90.000000
max         348.000000
Prompt length:  108
[2026-03-11 12:55:34,104] stat1010/INFO/locust.main: Run time limit set to 240 seconds
[2026-03-11 12:55:34,104] stat1010/INFO/locust.main: Starting Locust 2.31.5
[2026-03-11 12:55:34,104] stat1010/INFO/locust.runners: Ramping to 10 users at a rate of 10.00 per second
[2026-03-11 12:55:34,105] stat1010/INFO/locust.runners: All users spawned: {"Embeddings": 10} (10 total users)
[2026-03-11 12:59:33,599] stat1010/INFO/locust.main: --run-time limit reached, shutting down
Load test results are within the threshold
[2026-03-11 12:59:33,689] stat1010/INFO/locust.main: Shutting down (exit code 0)
Type     Name                                                                          # reqs      # fails |    Avg     Min     Max    Med |   req/s  failures/s
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
POST     /v1/models/qwen3-embedding:predict                                             23277     0(0.00%) |     51      11     651     51 |   97.17        0.00
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
         Aggregated                                                                     23277     0(0.00%) |     51      11     651     51 |   97.17        0.00

Response time percentiles (approximated)
Type     Name                                                                                  50%    66%    75%    80%    90%    95%    98%    99%  99.9% 99.99%   100% # reqs
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
POST     /v1/models/qwen3-embedding:predict                                                     51     56     60     62     68     73     80     86    290    620    650  23277
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
         Aggregated                                                                             51     56     60     62     68     73     80     86    290    620    650  23277

after aiter changes on staging:

(venv) ozge@stat1010:~/repos/wiki/gerrit/inference-services/test/locust$ MODEL=embeddings locust --host https://inference.svc.eqiad.wmnet:30443 --users 10 --run-time 4m
Min length: 10, Max length: 350
       question_length
count      4610.000000
mean         78.490456
std          37.470874
min          20.000000
25%          58.000000
50%          73.000000
75%          90.000000
max         348.000000
Prompt length:  108
[2026-03-12 17:09:49,113] stat1010/INFO/locust.main: Run time limit set to 240 seconds
[2026-03-12 17:09:49,113] stat1010/INFO/locust.main: Starting Locust 2.31.5
[2026-03-12 17:09:49,114] stat1010/INFO/locust.runners: Ramping to 10 users at a rate of 10.00 per second
[2026-03-12 17:09:49,114] stat1010/INFO/locust.runners: All users spawned: {"Embeddings": 10} (10 total users)
[2026-03-12 17:13:47,638] stat1010/INFO/locust.main: --run-time limit reached, shutting down
Load test results are within the threshold
[2026-03-12 17:13:47,741] stat1010/INFO/locust.main: Shutting down (exit code 0)
Type     Name                                                                          # reqs      # fails |    Avg     Min     Max    Med |   req/s  failures/s
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
POST     /v1/models/qwen3-embedding-staging:predict                                     23598     0(0.00%) |     49      12     894     48 |   98.90        0.00
--------|----------------------------------------------------------------------------|-------|-------------|-------|-------|-------|-------|--------|-----------
         Aggregated                                                                     23598     0(0.00%) |     49      12     894     48 |   98.90        0.00

Response time percentiles (approximated)
Type     Name                                                                                  50%    66%    75%    80%    90%    95%    98%    99%  99.9% 99.99%   100% # reqs
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
POST     /v1/models/qwen3-embedding-staging:predict                                             48     54     57     59     65     70     77     84    300    340    890  23598
--------|--------------------------------------------------------------------------------|--------|------|------|------|------|------|------|------|------|------|------|------
         Aggregated                                                                             48     54     57     59     65     70     77     84    300    340    890  23598

kevinbazira closed subtask T419650: Add ROCm build dependencies to wmf-debian-vllm image to support AITER kernel compilation as Resolved.Mar 13 2026, 5:51 AM

Thank you for running the load tests @OKarakaya-WMF. Below is a consolidated report of the optimizations we have implemented and their corresponding performance results for the embeddings isvc that is running on an MI300x GPU:

Inference Backend	Requests/s (Throughput)	Median Latency	p90	Failure Rate	Total Requests	Report	Status
HF Transformers+CK FlashAttention (T412338#11498835)	33.94	240ms	280ms	0.06%	8130	T418976#11676658	Stable. Was previously deployed in LiftWing prod and used for semantic search experiments.
vLLM+Triton FlashAttention (T418976#11676540)	94.38	52ms	71ms	0.00%	22610	T418976#11676658	Stable. Is currently deployed in LiftWing prod and used in semantic search prod.
vLLM+AITER (T418976#11690847)	98.90	48ms	65ms	0.00%	23598	T418976#11676658	Still under evaluation. We have experienced a couple of issues with: JIT compilation of the AITER kernels (P89805) and V1 engine prefix caching (P89849). Although these issues have been resolved, we'll continue monitoring this optimization to confirm it's stable before deploying it to LiftWing prod.