No performance difference with vLLM: qwen2.5-7b vs qwen2.5-7b-MLA

[nikita-yatchenko]

Hello,

I ran an experiment with two models: base qwen2.5-7b vs qwen2.5-7b-MLA and found that there was no difference in either speed nor KV Cache utilization.

qwen2.5-7b-MLA was converted using your script: bash scripts/convert/qwen2.5-7B-Instruct.sh
both models were hosted with vLLM.

Results for base qwen2.5-7b:

Results for base qwen2.5-7b-MLA:

Could you please specify how exactly you measured increase in output throughout and KV Cache saving?

[ItGirls]

Hello,

I ran an experiment with two models: base qwen2.5-7b vs qwen2.5-7b-MLA and found that there was no difference in either speed nor KV Cache utilization.

qwen2.5-7b-MLA was converted using your script: bash scripts/convert/qwen2.5-7B-Instruct.sh both models were hosted with vLLM.

Results for base qwen2.5-7b:

Results for base qwen2.5-7b-MLA:

Could you please specify how exactly you measured increase in output throughout and KV Cache saving?

Hello, what you use to run qwen2.5-7b-MLA? using Transformers.AutoModelForCausalLM or # using vllm.LLM?

[nikita-yatchenko]

vllm

linked issue:
https://github.com/fxmeng/TransMLA/issues/35

[Stanleytowne]

Hi there,

Could you please provide the inference code and environment details you used for your test?

The KV cache compression ratio can be calculated theoretically (as explained in Issue #36). For the Qwen2.5-7B model, the original KV cache size per token/layer is 2 * 4 * 128 = 1024. The configuration in our provided script is 512 + 64, which results in a theoretical KV cache reduction of -43.75%.

Also, we will be uploading our throughput testing script shortly.

Thanks!

[nikita-yatchenko]

Thank you for the response

here is docker-compose.yml that i used:

services:
  vllm-server-original:
    build:
      context: .
      dockerfile: Dockerfile
    container_name: vllm-server-original
    runtime: nvidia
    ipc: host
    ports:
      - "8080:8080"
    volumes:
      - .:/usr/src
    shm_size: '8gb'
    environment:
      - DO_NOT_TRACK=1
      - NVIDIA_VISIBLE_DEVICES=all
      - NCCL_DEBUG=INFO
      - NCCL_NVLS_ENABLE=0
      - CUDA_VISIBLE_DEVICES=0
      - CUDA_LAUNCH_BLOCKING=1
      - VLLM_USE_PRECOMPUTED_WEIGHTS=true
    command: >
      vllm serve Qwen/Qwen2.5-3B-Instruct
      --port 8080
      --gpu_memory_utilization 0.95
      --served-model-name Qwen/Qwen2.5-3B-Instruct
      --max_model_len 32768
    networks:
      - default
    
  vllm-server-transmla:
    build:
      context: .
      dockerfile: Dockerfile
    container_name: vllm-server-transmla
    runtime: nvidia
    ipc: host
    ports:
      - "8081:8080"
    volumes:
      - .:/usr/src
      - /var/llm-evaluation:/var/llm-evaluation
    shm_size: '8gb'
    environment:
      - DO_NOT_TRACK=1
      - NVIDIA_VISIBLE_DEVICES=all
      - PYTHONPATH=./TransMLA:${PYTHONPATH}
      - NCCL_DEBUG=INFO
      - NCCL_NVLS_ENABLE=0
      - CUDA_VISIBLE_DEVICES=1
      - CUDA_LAUNCH_BLOCKING=1
      - VLLM_USE_PRECOMPUTED_WEIGHTS=true
    command: >
      sh -c "
      cd /usr/src &&
      PYTHONPATH=/usr/src/TransMLA:\$PYTHONPATH \
      python launch_vllm_transmla_server.py \
      --model /var/llm-evaluation/qwen2_5-3B-Instruct-MLA-576 \
      --host 0.0.0.0 \
      --port 8080 \
      --served-model-name Qwen/Qwen2.5-3B-Instruct-MLA \
      --gpu-memory-utilization 0.95 \
      --max-model-len 32768 \
      --trust-remote-code \
      --dtype auto
      "

I used default script provided in your documentation:
qwen2.5-7B-Instruct.sh

but changed from 7B to 3B model

Inference:
Async requests with rps of 1. Prompts are different from requerst to request, but the same for each model

[Stanleytowne]

Thanks for providing your setup details.

This explains the behavior you're seeing. For the Qwen2.5-3B model, the parameters are head_dim=128 and num_kv_heads=2. Therefore, the original KV cache size is only 2 * 2 * 128 = 512.

Our script's setting is 512 + 64 = 576, which means that for this specific model, the KV cache size actually increases.

Could you please try again with the Qwen2.5-7B model? We have also uploaded our throughput testing code, with which we are able to get a significant speedup.

[nikita-yatchenko]

Just finished re training the model and tested with vLLM:

here are the logs for vLLM startup with 7B CONVERTED model:

INFO 09-23 13:36:03 [__init__.py:239] Automatically detected platform cuda.
Final args: {'host': '0.0.0.0', 'port': 8080, 'uvicorn_log_level': 'info', 'disable_uvicorn_access_log': False, 'allow_credentials': False, 'allowed_origins': ['*'], 'allowed_methods': ['*'], 'allowed_headers': ['*'], 'api_key': None, 'lora_modules': None, 'prompt_adapters': None, 'chat_template': None, 'chat_template_content_format': 'auto', 'response_role': 'assistant', 'ssl_keyfile': None, 'ssl_certfile': None, 'ssl_ca_certs': None, 'enable_ssl_refresh': False, 'ssl_cert_reqs': 0, 'root_path': None, 'middleware': [], 'return_tokens_as_token_ids': False, 'disable_frontend_multiprocessing': False, 'enable_request_id_headers': False, 'enable_auto_tool_choice': False, 'tool_call_parser': None, 'tool_parser_plugin': '', 'model': '/var/llm-evaluation/qwen2_5-7B-Instruct-response', 'task': 'auto', 'tokenizer': None, 'hf_config_path': None, 'skip_tokenizer_init': False, 'revision': None, 'code_revision': None, 'tokenizer_revision': None, 'tokenizer_mode': 'auto', 'trust_remote_code': True, 'allowed_local_media_path': None, 'load_format': 'auto', 'download_dir': None, 'model_loader_extra_config': None, 'use_tqdm_on_load': True, 'config_format': <ConfigFormat.AUTO: 'auto'>, 'dtype': 'auto', 'kv_cache_dtype': 'auto', 'max_model_len': 32768, 'guided_decoding_backend': 'auto', 'logits_processor_pattern': None, 'model_impl': 'auto', 'distributed_executor_backend': None, 'pipeline_parallel_size': 1, 'tensor_parallel_size': 1, 'data_parallel_size': 1, 'enable_expert_parallel': False, 'max_parallel_loading_workers': None, 'ray_workers_use_nsight': False, 'disable_custom_all_reduce': False, 'block_size': None, 'enable_prefix_caching': None, 'prefix_caching_hash_algo': 'builtin', 'disable_sliding_window': False, 'use_v2_block_manager': True, 'num_lookahead_slots': 0, 'seed': None, 'swap_space': 4, 'cpu_offload_gb': 0, 'gpu_memory_utilization': 0.95, 'num_gpu_blocks_override': None, 'max_num_batched_tokens': None, 'max_num_partial_prefills': 1, 'max_long_partial_prefills': 1, 'long_prefill_token_threshold': 0, 'max_num_seqs': None, 'max_logprobs': 20, 'disable_log_stats': False, 'quantization': None, 'rope_scaling': None, 'rope_theta': None, 'hf_token': None, 'hf_overrides': None, 'enforce_eager': False, 'max_seq_len_to_capture': 8192, 'tokenizer_pool_size': 0, 'tokenizer_pool_type': 'ray', 'tokenizer_pool_extra_config': None, 'limit_mm_per_prompt': None, 'mm_processor_kwargs': None, 'disable_mm_preprocessor_cache': False, 'enable_lora': False, 'enable_lora_bias': False, 'max_loras': 1, 'max_lora_rank': 16, 'lora_extra_vocab_size': 256, 'lora_dtype': 'auto', 'long_lora_scaling_factors': None, 'max_cpu_loras': None, 'fully_sharded_loras': False, 'enable_prompt_adapter': False, 'max_prompt_adapters': 1, 'max_prompt_adapter_token': 0, 'device': 'auto', 'num_scheduler_steps': 1, 'multi_step_stream_outputs': True, 'scheduler_delay_factor': 0.0, 'enable_chunked_prefill': None, 'speculative_config': None, 'ignore_patterns': [], 'preemption_mode': None, 'served_model_name': ['Qwen/Qwen2.5-7B-Instruct-MLA'], 'qlora_adapter_name_or_path': None, 'show_hidden_metrics_for_version': None, 'otlp_traces_endpoint': None, 'collect_detailed_traces': None, 'disable_async_output_proc': False, 'scheduling_policy': 'fcfs', 'scheduler_cls': 'vllm.core.scheduler.Scheduler', 'override_neuron_config': None, 'override_pooler_config': None, 'compilation_config': None, 'kv_transfer_config': None, 'worker_cls': 'auto', 'worker_extension_cls': '', 'generation_config': 'auto', 'override_generation_config': None, 'enable_sleep_mode': False, 'calculate_kv_scales': False, 'additional_config': None, 'enable_reasoning': False, 'reasoning_parser': None, 'disable_cascade_attn': False, 'disable_chunked_mm_input': False, 'disable_log_requests': False, 'max_log_len': None, 'disable_fastapi_docs': False, 'enable_prompt_tokens_details': False, 'enable_server_load_tracking': False}
INFO 09-23 13:36:06 [api_server.py:1034] vLLM API server version 0.8.4
INFO 09-23 13:36:06 [api_server.py:1035] args: Namespace(host='0.0.0.0', port=8080, uvicorn_log_level='info', disable_uvicorn_access_log=False, allow_credentials=False, allowed_origins=['*'], allowed_methods=['*'], allowed_headers=['*'], api_key=None, lora_modules=None, prompt_adapters=None, chat_template=None, chat_template_content_format='auto', response_role='assistant', ssl_keyfile=None, ssl_certfile=None, ssl_ca_certs=None, enable_ssl_refresh=False, ssl_cert_reqs=0, root_path=None, middleware=[], return_tokens_as_token_ids=False, disable_frontend_multiprocessing=False, enable_request_id_headers=False, enable_auto_tool_choice=False, tool_call_parser=None, tool_parser_plugin='', model='**/var/llm-evaluation/qwen2_5-7B-Instruct-response**', task='auto', tokenizer=None, hf_config_path=None, skip_tokenizer_init=False, revision=None, code_revision=None, tokenizer_revision=None, tokenizer_mode='auto', trust_remote_code=True, allowed_local_media_path=None, load_format='auto', download_dir=None, model_loader_extra_config=None, use_tqdm_on_load=True, config_format=<ConfigFormat.AUTO: 'auto'>, dtype='auto', kv_cache_dtype='auto', max_model_len=32768, guided_decoding_backend='auto', logits_processor_pattern=None, model_impl='auto', distributed_executor_backend=None, pipeline_parallel_size=1, tensor_parallel_size=1, data_parallel_size=1, enable_expert_parallel=False, max_parallel_loading_workers=None, ray_workers_use_nsight=False, disable_custom_all_reduce=False, block_size=None, enable_prefix_caching=None, prefix_caching_hash_algo='builtin', disable_sliding_window=False, use_v2_block_manager=True, num_lookahead_slots=0, seed=None, swap_space=4, cpu_offload_gb=0, gpu_memory_utilization=0.95, num_gpu_blocks_override=None, max_num_batched_tokens=None, max_num_partial_prefills=1, max_long_partial_prefills=1, long_prefill_token_threshold=0, max_num_seqs=None, max_logprobs=20, disable_log_stats=False, quantization=None, rope_scaling=None, rope_theta=None, hf_token=None, hf_overrides=None, enforce_eager=False, max_seq_len_to_capture=8192, tokenizer_pool_size=0, tokenizer_pool_type='ray', tokenizer_pool_extra_config=None, limit_mm_per_prompt=None, mm_processor_kwargs=None, disable_mm_preprocessor_cache=False, enable_lora=False, enable_lora_bias=False, max_loras=1, max_lora_rank=16, lora_extra_vocab_size=256, lora_dtype='auto', long_lora_scaling_factors=None, max_cpu_loras=None, fully_sharded_loras=False, enable_prompt_adapter=False, max_prompt_adapters=1, max_prompt_adapter_token=0, device='auto', num_scheduler_steps=1, multi_step_stream_outputs=True, scheduler_delay_factor=0.0, enable_chunked_prefill=None, speculative_config=None, ignore_patterns=[], preemption_mode=None, served_model_name=['Qwen/Qwen2.5-7B-Instruct-MLA'], qlora_adapter_name_or_path=None, show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, disable_async_output_proc=False, scheduling_policy='fcfs', scheduler_cls='vllm.core.scheduler.Scheduler', override_neuron_config=None, override_pooler_config=None, compilation_config=None, kv_transfer_config=None, worker_cls='auto', worker_extension_cls='', generation_config='auto', override_generation_config=None, enable_sleep_mode=False, calculate_kv_scales=False, additional_config=None, enable_reasoning=False, reasoning_parser=None, disable_cascade_attn=False, disable_chunked_mm_input=False, disable_log_requests=False, max_log_len=None, disable_fastapi_docs=False, enable_prompt_tokens_details=False, enable_server_load_tracking=False)
INFO 09-23 13:36:06 [config.py:2832] Downcasting torch.float32 to torch.float16.
INFO 09-23 13:36:06 [config.py:689] This model supports multiple tasks: {'generate', 'embed', 'score', 'classify', 'reward'}. Defaulting to 'generate'.
INFO 09-23 13:36:06 [config.py:1901] Chunked prefill is enabled with max_num_batched_tokens=2048.
INFO 09-23 13:36:11 [__init__.py:239] Automatically detected platform cuda.
INFO 09-23 13:36:13 [core.py:61] Initializing a V1 LLM engine (v0.8.4) with config: model='/var/llm-evaluation/qwen2_5-7B-Instruct-response', speculative_config=None, tokenizer='/var/llm-evaluation/qwen2_5-7B-Instruct-response', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.float16, max_seq_len=32768, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='auto', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=None, served_model_name=Qwen/Qwen2.5-7B-Instruct-MLA, num_scheduler_steps=1, mulLoading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:01<00:04,  1.51s/it]
_prefill_enabled=True, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"level":3,"custom_ops":["none"],"splitting_ops":["vllm.unified_attention","vllm.unified_attention_with_output"],"use_inductor":true,"compile_sizes":[],"use_cudagraph":true,"cudagraph_num_of_warmups":1,"cudagraph_capture_sizes":[512,504,496,488,480,472,464,456,448,440,432,424,416,408,400,392,384,376,368,360,352,344,336,328,320,312,304,296,288,280,272,264,256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":512}
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:03<00:03,  1.62s/it]
WARNING 09-23 13:36:14 [utils.py:2444] Methods determine_num_available_blocks,device_config,get_cache_block_size_bytes,initialize_cache not implemented in <vllm.v1.worker.gpu_worker.Worker object at 0x7fe38456ed80>
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:04<00:01,  1.68s/it]
INFO 09-23 13:36:14 [parallel_state.py:959] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0
INFO 09-23 13:36:14 [cuda.py:187] Using Triton MLA backend on V1 engine.
WARNING 09-23 13:36:14 [triton_decode_attention.py:44] The following error message 'operation scheduled before its operands' can be ignored.
INFO 09-23 13:36:14 [gpu_model_runner.py:1276] Starting to load model /var/llm-evaluation/qwen2_5-7B-Instruct-response...
WARNING 09-23 13:36:14 [topk_topp_sampler.py:69] FlashInfer is not available. Falling back to the PyTorch-native implementation of top-p & top-k sampling. For the best performance, please install FlashInfer.
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:05<00:00,  1.18s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:05<00:00,  1.34s/it]

INFO 09-23 13:36:20 [loader.py:458] Loading weights took 5.38 seconds
INFO 09-23 13:36:20 [gpu_model_runner.py:1291] Model loading took 14.6750 GiB and 5.568219 seconds
INFO 09-23 13:36:25 [backends.py:416] Using cache directory: /root/.cache/vllm/torch_compile_cache/3fbb3796e0/rank_0_0 for vLLM's torch.compile
INFO 09-23 13:36:25 [backends.py:426] Dynamo bytecode transform time: 4.73 s
INFO 09-23 13:36:27 [backends.py:132] Cache the graph of shape None for later use
INFO 09-23 13:36:42 [backends.py:144] Compiling a graph for general shape takes 16.87 s
INFO 09-23 13:36:49 [monitor.py:33] torch.compile takes 21.59 s in total
INFO 09-23 13:36:51 [kv_cache_utils.py:634] GPU KV cache size: 1,654,224 tokens
INFO 09-23 13:36:51 [kv_cache_utils.py:637] Maximum concurrency for 32,768 tokens per request: 50.48x
INFO 09-23 13:37:13 [gpu_model_runner.py:1626] Graph capturing finished in 23 secs, took 1.24 GiB
INFO 09-23 13:37:13 [core.py:163] init engine (profile, create kv cache, warmup model) took 53.47 seconds
INFO 09-23 13:37:14 [core_client.py:435] Core engine process 0 ready.
WARNING 09-23 13:37:14 [config.py:1177] Default sampling parameters have been overridden by the model's Hugging Face generation config recommended from the model creator. If this is not intended, please relaunch vLLM instance with `--generation-config vllm`.
INFO 09-23 13:37:14 [serving_chat.py:118] Using default chat sampling params from model: {'repetition_penalty': 1.05, 'temperature': 0.7, 'top_k': 20, 'top_p': 0.8}
INFO 09-23 13:37:14 [serving_completion.py:61] Using default completion sampling params from model: {'repetition_penalty': 1.05, 'temperature': 0.7, 'top_k': 20, 'top_p': 0.8}
INFO 09-23 13:37:14 [api_server.py:1081] Starting vLLM API server on http://0.0.0.0:8080
INFO 09-23 13:37:14 [launcher.py:26] Available routes are:
INFO 09-23 13:37:14 [launcher.py:34] Route: /openapi.json, Methods: GET, HEAD
INFO 09-23 13:37:14 [launcher.py:34] Route: /docs, Methods: GET, HEAD
INFO 09-23 13:37:14 [launcher.py:34] Route: /docs/oauth2-redirect, Methods: GET, HEAD
INFO 09-23 13:37:14 [launcher.py:34] Route: /redoc, Methods: GET, HEAD
INFO 09-23 13:37:14 [launcher.py:34] Route: /health, Methods: GET
INFO 09-23 13:37:14 [launcher.py:34] Route: /load, Methods: GET
INFO 09-23 13:37:14 [launcher.py:34] Route: /ping, Methods: GET, POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /tokenize, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /detokenize, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/models, Methods: GET
INFO 09-23 13:37:14 [launcher.py:34] Route: /version, Methods: GET
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/chat/completions, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/completions, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/embeddings, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /pooling, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /score, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/score, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/audio/transcriptions, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /rerank, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /v1/rerank, Methods: POST
INFO:     Started server process [7]
INFO 09-23 13:37:14 [launcher.py:34] Route: /v2/rerank, Methods: POST
INFO 09-23 13:37:14 [launcher.py:34] Route: /invocations, Methods: POST
INFO:     Waiting for application startup.
INFO:     Application startup complete.

qwen2.5-7B-Instruct.sh:

model_path=Qwen/Qwen2.5-7B-Instruct
save_path=outputs/qwen2_5-7B-Instruct-response
eval_batch_size=8

python transmla/converter.py \
    --model-path $model_path \
    --save-path $save_path \
    --freqfold 4 \
    --ppl-eval-batch-size $eval_batch_size \
    --kv-lora-rank 512 \
    --qk-mqa-dim 64

Here is the speedtest summary for ORIGINAL 7B model:

============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  1030.60   
Total input tokens:                      4493235   
Total generated tokens:                  1001000   
Request throughput (req/s):              0.97  
Output token throughput (tok/s):         971.28   
Total Token throughput (tok/s):          5331.09   
---------------Time to First Token----------------
Mean TTFT (ms):                          355.05    
Median TTFT (ms):                        405.27    
P99 TTFT (ms):                           463.70    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          44.32     
Median TPOT (ms):                        45.36     
P99 TPOT (ms):                           49.19     
---------------Inter-token Latency----------------
Mean ITL (ms):                           22480.88  
Median ITL (ms):                         22405.70  
P99 ITL (ms):                            46747.72  
==================================================

Here is the speedtest summary for CONVERTED 7B model (MLA):

============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  1142.30
Total input tokens:                      4493235   
Total generated tokens:                  1001000   
Request throughput (req/s):              0.88
Output token throughput (tok/s):         876.30
Total Token throughput (tok/s):          4809.79   
---------------Time to First Token----------------
Mean TTFT (ms):                          16804.21  
Median TTFT (ms):                        12832.37  
P99 TTFT (ms):                           43750.08  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          238.39    
Median TPOT (ms):                        255.27    
P99 TPOT (ms):                           286.48    
---------------Inter-token Latency----------------
Mean ITL (ms):                           137518.04 
Median ITL (ms):                         134811.90 
P99 ITL (ms):                            284655.30 
==================================================

It is visible slower

The picture for the GPU usage is similar to the one above (i.e. I cannot observe reduced GPU memory usage)

Mean input token length:
4568.6

Fixed output token length (EOS ignore):
1000

[nikita-yatchenko]

I also ran your throughput script with Qwen2.5-7b

WARNING 09-25 12:27:49 [topk_topp_sampler.py:69] FlashInfer is not available. Falling back to the PyTorch-native implementation of top-p & top-k sampling. For the best performance, please install FlashInfer.
INFO 09-25 12:27:50 [weight_utils.py:265] Using model weights format ['*.safetensors']
Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:01<00:05,  1.95s/it]
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:06<00:06,  3.47s/it]
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:11<00:04,  4.10s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:13<00:00,  3.37s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:13<00:00,  3.40s/it]

INFO 09-25 12:28:04 [loader.py:458] Loading weights took 13.70 seconds
INFO 09-25 12:28:04 [gpu_model_runner.py:1291] Model loading took 14.2487 GiB and 14.722004 seconds
INFO 09-25 12:28:12 [backends.py:416] Using cache directory: /root/.cache/vllm/torch_compile_cache/e76f81c887/rank_0_0 for vLLM's torch.compile
INFO 09-25 12:28:12 [backends.py:426] Dynamo bytecode transform time: 7.64 s
INFO 09-25 12:28:12 [backends.py:115] Directly load the compiled graph for shape None from the cache
INFO 09-25 12:28:21 [monitor.py:33] torch.compile takes 7.64 s in total
INFO 09-25 12:28:22 [kv_cache_utils.py:634] GPU KV cache size: 1,029,600 tokens
INFO 09-25 12:28:22 [kv_cache_utils.py:637] Maximum concurrency for 32,768 tokens per request: 31.42x
INFO 09-25 12:28:49 [gpu_model_runner.py:1626] Graph capturing finished in 26 secs, took 1.45 GiB
INFO 09-25 12:28:49 [core.py:163] init engine (profile, create kv cache, warmup model) took 44.81 seconds
INFO 09-25 12:28:49 [core_client.py:435] Core engine process 0 ready.
Processed prompts: 100%|█████████████████████████████████████| 100/100 [41:56<00:00, 25.17s/it, est. speed input: 650.97 toks/s, output: 650.97 toks/s]
Throughput: 0.04 requests/s, 1299.14 total tokens/s, 649.57 output tokens/s
Total num prompt tokens:  1638400
Total num output tokens:  1638400
Running /app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response with input length 16384 and output length 16384
INFO 09-25 13:11:02 [__init__.py:239] Automatically detected platform cuda.
When dataset path is not set, it will default to random dataset
Namespace(backend='vllm', dataset_name='random', dataset=None, dataset_path=None, input_len=16384, output_len=16384, n=1, num_prompts=100, hf_max_batch_size=None, output_json=None, async_engine=False, disable_frontend_multiprocessing=False, disable_detokenize=False, lora_path=None, prefix_len=0, random_range_ratio=0.0, hf_subset=None, hf_split=None, model='/app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response', task='auto', tokenizer='/app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response', hf_config_path=None, skip_tokenizer_init=False, revision=None, code_revision=None, tokenizer_revision=None, tokenizer_mode='auto', trust_remote_code=False, allowed_local_media_path=None, load_format='auto', download_dir=None, model_loader_extra_config=None, use_tqdm_on_load=True, config_format=<ConfigFormat.AUTO: 'auto'>, dtype='auto', kv_cache_dtype='auto', max_model_len=None, guided_decoding_backend='auto', logits_processor_pattern=None, model_impl='auto', distributed_executor_backend=None, pipeline_parallel_size=1, tensor_parallel_size=1, data_parallel_size=1, enable_expert_parallel=False, max_parallel_loading_workers=None, ray_workers_use_nsight=False, disable_custom_all_reduce=False, block_size=None, enable_prefix_caching=None, prefix_caching_hash_algo='builtin', disable_sliding_window=False, use_v2_block_manager=True, num_lookahead_slots=0, seed=0, swap_space=4, cpu_offload_gb=0, gpu_memory_utilization=0.9, num_gpu_blocks_override=None, max_num_batched_tokens=None, max_num_partial_prefills=1, max_long_partial_prefills=1, long_prefill_token_threshold=0, max_num_seqs=None, max_logprobs=20, disable_log_stats=False, quantization=None, rope_scaling=None, rope_theta=None, hf_token=None, hf_overrides=None, enforce_eager=False, max_seq_len_to_capture=8192, tokenizer_pool_size=0, tokenizer_pool_type='ray', tokenizer_pool_extra_config=None, limit_mm_per_prompt=None, mm_processor_kwargs=None, disable_mm_preprocessor_cache=False, enable_lora=False, enable_lora_bias=False, max_loras=1, max_lora_rank=16, lora_extra_vocab_size=256, lora_dtype='auto', long_lora_scaling_factors=None, max_cpu_loras=None, fully_sharded_loras=False, enable_prompt_adapter=False, max_prompt_adapters=1, max_prompt_adapter_token=0, device='auto', num_scheduler_steps=1, multi_step_stream_outputs=True, scheduler_delay_factor=0.0, enable_chunked_prefill=None, speculative_config=None, ignore_patterns=[], preemption_mode=None, served_model_name=None, qlora_adapter_name_or_path=None, show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, disable_async_output_proc=False, scheduling_policy='fcfs', scheduler_cls='vllm.core.scheduler.Scheduler', override_neuron_config=None, override_pooler_config=None, compilation_config=None, kv_transfer_config=None, worker_cls='auto', worker_extension_cls='', generation_config='auto', override_generation_config=None, enable_sleep_mode=False, calculate_kv_scales=False, additional_config=None, enable_reasoning=False, reasoning_parser=None, disable_cascade_attn=False, disable_chunked_mm_input=False, disable_log_requests=False)
INFO 09-25 13:11:14 [config.py:2832] Downcasting torch.float32 to torch.float16.
INFO 09-25 13:11:14 [config.py:689] This model supports multiple tasks: {'reward', 'embed', 'score', 'classify', 'generate'}. Defaulting to 'generate'.
INFO 09-25 13:11:14 [config.py:1901] Chunked prefill is enabled with max_num_batched_tokens=8192.
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
INFO 09-25 13:11:14 [core.py:61] Initializing a V1 LLM engine (v0.8.4) with config: model='/app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response', speculative_config=None, tokenizer='/app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=32768, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='auto', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=/app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=True, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"level":3,"custom_ops":["none"],"splitting_ops":["vllm.unified_attention","vllm.unified_attention_with_output"],"use_inductor":true,"compile_sizes":[],"use_cudagraph":true,"cudagraph_num_of_warmups":1,"cudagraph_capture_sizes":[512,504,496,488,480,472,464,456,448,440,432,424,416,408,400,392,384,376,368,360,352,344,336,328,320,312,304,296,288,280,272,264,256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":512}
WARNING 09-25 13:11:14 [utils.py:2444] Methods determine_num_available_blocks,device_config,get_cache_block_size_bytes,initialize_cache not implemented in <vllm.v1.worker.gpu_worker.Worker object at 0x7f14b65d64b0>
INFO 09-25 13:11:15 [parallel_state.py:959] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0
INFO 09-25 13:11:15 [cuda.py:187] Using Triton MLA backend on V1 engine.
WARNING 09-25 13:11:15 [triton_decode_attention.py:44] The following error message 'operation scheduled before its operands' can be ignored.
INFO 09-25 13:11:15 [gpu_model_runner.py:1276] Starting to load model /app/projects/adaptation_qga_mla/custom_model_vllm/TransMLA/outputs/qwen2_5-7B-Instruct-response...
WARNING 09-25 13:11:15 [topk_topp_sampler.py:69] FlashInfer is not available. Falling back to the PyTorch-native implementation of top-p & top-k sampling. For the best performance, please install FlashInfer.
Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:05<00:17,  5.77s/it]
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:12<00:12,  6.26s/it]
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:18<00:06,  6.41s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:20<00:00,  4.47s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:20<00:00,  5.12s/it]

INFO 09-25 13:11:36 [loader.py:458] Loading weights took 20.49 seconds
INFO 09-25 13:11:36 [gpu_model_runner.py:1291] Model loading took 14.6739 GiB and 20.695562 seconds
INFO 09-25 13:11:41 [backends.py:416] Using cache directory: /root/.cache/vllm/torch_compile_cache/9ad4f7b13d/rank_0_0 for vLLM's torch.compile
INFO 09-25 13:11:41 [backends.py:426] Dynamo bytecode transform time: 4.33 s
INFO 09-25 13:11:41 [backends.py:115] Directly load the compiled graph for shape None from the cache
INFO 09-25 13:11:44 [monitor.py:33] torch.compile takes 4.33 s in total
INFO 09-25 13:11:45 [kv_cache_utils.py:634] GPU KV cache size: 1,808,608 tokens
INFO 09-25 13:11:45 [kv_cache_utils.py:637] Maximum concurrency for 32,768 tokens per request: 55.19x
INFO 09-25 13:12:08 [gpu_model_runner.py:1626] Graph capturing finished in 23 secs, took 1.24 GiB
INFO 09-25 13:12:08 [core.py:163] init engine (profile, create kv cache, warmup model) took 31.98 seconds
INFO 09-25 13:12:08 [core_client.py:435] Core engine process 0 ready.
Processed prompts: 100%|███████████████████████████████████| 100/100 [1:05:17<00:00, 39.17s/it, est. speed input: 418.25 toks/s, output: 418.25 toks/s]
Throughput: 0.03 requests/s, 835.18 total tokens/s, 417.59 output tokens/s
Total num prompt tokens:  1638400
Total num output tokens:  1638400

as you can see there is a slow down when using qwen2.5 MLA

However, when I repeated your Llama2-7B experiment I confirm that there is twofold increase in throughput..

Perhaps current approach is not optimized for Qwen?

I would very much appreciate if you ran conversion for Qwen2.5-7B and throughput test - just to compare result.

Additionally, I would like to point out that Llama2-7B is basically an MHA, not GQA. GQA is used in the 70B model

[xlim1997]

Qwen-2.5 already keeps the KV cache small via GQA. For the 7B model, num_attention_heads=28 and num_kv_heads=4, meaning one KV group is shared by 7 attention heads. Under this setup, to obtain additional KV-memory savings from MLA you have to push the latent rank very low; otherwise the marginal reduction over the “4 KV groups” baseline is limited. At the same time, MLA requires an up-projection from the compact latent K/V back to the full dimension at every step of autoregressive decoding, which introduces non-trivial compute overhead; with long KV sequences this overhead scales roughly linearly with sequence length and can dominate latency.

[Wfd567]

@nikita-yatchenko
python launch_vllm_transmla_server.py
--model /var/llm-evaluation/qwen2_5-3B-Instruct-MLA-576
--host 0.0.0.0
--port 8080
--served-model-name Qwen/Qwen2.5-3B-Instruct-MLA
--gpu-memory-utilization 0.95
Hello, I encountered a small issue when using vllm for online deployment. Could you share the launch_vllm_transmla_server.py script?I am looking forward to your early reply.
All the best