[Attention] Flash MLA for V1

[LucasWilkinson]

use via: VLLM_ATTENTION_BACKEND=FLASHMLA VLLM_USE_V1=1

Results:

https://docs.google.com/spreadsheets/d/1toxQVaA7UPhmY57kv08Wdq0xtIcU8oBbqb9RE3Wy2_E/edit?usp=sharing

[LiuXiaoxuanPKU]

QQ: Just curiosity, what's the main reason that V1 is slower than V0 (say all use Flash MLA, and we look at ITL)? Is it because of chunked prefill?

[LucasWilkinson]

QQ: Just curiosity, what's the main reason that V1 is slower than V0 (say all use Flash MLA, and we look at ITL)? Is it because of chunked prefill?

Its mostly likely because in V1 CUDA graphs are not used for attention, we need to keep optimizing away alot of the small operations in MLA (for low QPS, for the throughput case it may be chunked prefill)

[mgoin]

This looks clean to me, nice work! Have you run an accuracy smoke test?

[LucasWilkinson]

This looks clean to me, nice work! Have you run an accuracy smoke test?

VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHMLA lm_eval --model vllm --model_args pretrained=deepseek-ai/DeepSeek-V2-Lite-Chat,tensor_parallel_size=2,dtype=auto,gpu_memory_utilization=0.9,trust_remote_code=True,max_model_len=16384,enforce_eager=True --task gsm8k --num_fewshot=5 --limit 100

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.66|±  |0.0476|
|     |       |strict-match    |     5|exact_match|↑  | 0.66|±  |0.0476|

[samuellees]

use via: VLLM_ATTENTION_BACKEND=FLASHMLA VLLM_USE_V1=1

Results:

https://docs.google.com/spreadsheets/d/1toxQVaA7UPhmY57kv08Wdq0xtIcU8oBbqb9RE3Wy2_E/edit?usp=sharing

@LucasWilkinson Hello! Which GPU hardware did these tests run on? It seems that out-of-memory (OOM) errors occur on most NVIDIA GPUs when performing the "Long context" test.

[zuozi2810]

Hi, @LucasWilkinson, I got errors while running deepseek r1 awq with FlashMLA and V1 engine in 8 * H100:

command:
VLLM_USE_V1=1 VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_ATTENTION_BACKEND=FLASHMLA python benchmark_throughput.py --model /workspace/model --trust-remote-code --input-len 2000 --output-len 1000 --num-prompts 100 -tp 8 --gpu_memory_utilization 0.97 --max-model-len 5120

[LucasWilkinson]

Hi, @LucasWilkinson, I got errors while running deepseek r1 awq with FlashMLA and V1 engine in 8 * H100:

command: VLLM_USE_V1=1 VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_ATTENTION_BACKEND=FLASHMLA python benchmark_throughput.py --model /workspace/model --trust-remote-code --input-len 2000 --output-len 1000 --num-prompts 100 -tp 8 --gpu_memory_utilization 0.97 --max-model-len 5120

@zuozi2810 can you please provide the full log and hugging face link to a failing model, that would be really helpful in debugging

[LucasWilkinson]

use via: VLLM_ATTENTION_BACKEND=FLASHMLA VLLM_USE_V1=1
Results:
https://docs.google.com/spreadsheets/d/1toxQVaA7UPhmY57kv08Wdq0xtIcU8oBbqb9RE3Wy2_E/edit?usp=sharing

@LucasWilkinson Hello! Which GPU hardware did these tests run on? It seems that out-of-memory (OOM) errors occur on most NVIDIA GPUs when performing the "Long context" test.

8xH200, what setup are you seeing OOM on?

[samuellees]

use via: VLLM_ATTENTION_BACKEND=FLASHMLA VLLM_USE_V1=1
Results:
https://docs.google.com/spreadsheets/d/1toxQVaA7UPhmY57kv08Wdq0xtIcU8oBbqb9RE3Wy2_E/edit?usp=sharing

@LucasWilkinson Hello! Which GPU hardware did these tests run on? It seems that out-of-memory (OOM) errors occur on most NVIDIA GPUs when performing the "Long context" test.

8xH200, what setup are you seeing OOM on?

8xH200 have enough GPU memory. My setup is 8xH20. Thank you for your reply.