[Attention] MLA get rid of materialization

[LucasWilkinson]

Based on these calculations:

https://docs.google.com/spreadsheets/d/17eoqEbhblvtNsRRlFSjCQnEXZiBxtLgZGKD4IgZUz38/edit?usp=sharing

It's actually better to just not materialize the absorbed W_Q_UK and W_UV_O as it reduces memory usage (and total flops) and instead compute using sequential matmuls. One issue is that we do not have an FP8 bmm (which is needed if not materializing the absorbed matrix, materializing absorbing allowed us to bypass this), so we instead decompress the matrices involved in the bmm to fp16/bf16. This also has the added benefit of dramatically reducing complexity.

This PR is needed for DP attention since without it the weight materialization eats up too much of the GPU memory to make DP beneficial.

This PR (minor regression in short context but seems worth it given the saved memory boosts long-context and enables DP-attention, also the short context measurements are a bit noisy)

  backend  input_tokens  output_tokens  output_toks/s     req/s  median_itl_ms  median_ttft_ms
3    vllm          1000           1000    1323.397915  1.323398      29.755860     2307.676603
2    vllm          5000           1000    1041.455043  1.041455      33.205620     5491.457423
4    vllm         10000           1000     874.563079  0.874563      36.871404     8508.498527
1    vllm         32000           1000     190.195698  0.190196      35.948243   108055.433401

Baseline (#14769)

  backend  input_tokens  output_tokens  output_toks/s     req/s  median_itl_ms  median_ttft_ms
3    vllm          1000           1000    1380.973383  1.380973      30.029079     2223.776098
2    vllm          5000           1000    1047.717832  1.047718      33.303024     5499.557093
4    vllm         10000           1000     586.460476  0.586460      36.855329     8512.637936
1    vllm         32000           1000     162.816157  0.162816      42.981104   115515.806204

correctness tests:

VLLM_USE_V1=1  vllm serve /home/vllm-dev/DeepSeek-R1 --tensor-parallel-size 8 --trust-remote-code --disable-log-requests

lm_eval --model local-completions --tasks gsm8k --model_args model=/home/vllm-dev/DeepSeek-R1,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=5,max_retries=3,tokenized_request
s=False --limit 100

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

VLLM_USE_V1=0 vllm serve /home/vllm-dev/DeepSeek-R1 --tensor-parallel-size 8 --trust-remote-code --disable-log-requests

lm_eval --model local-completions --tasks gsm8k --model_args model=/home/vllm-dev/DeepSeek-R1,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=5,max_retries=3,tokenized_request
s=False --limit 100

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

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[simon-mo]

Make sense for the memory savings

[erictanjn]

Based on these calculations:

https://docs.google.com/spreadsheets/d/17eoqEbhblvtNsRRlFSjCQnEXZiBxtLgZGKD4IgZUz38/edit?usp=sharing

It's actually better to just not materialize the absorbed W_Q_UK and W_UV_O as it reduces memory usage (and total flops) and instead compute using sequential matmuls. One issue is that we do not have an FP8 bmm (which is needed if not materializing the absorbed matrix, materializing absorbing allowed us to bypass this), so we instead decompress the matrices involved in the bmm to fp16/bf16. This also has the added benefit of dramatically reducing complexity.

This PR is needed for DP attention since without it the weight materialization eats up too much of the GPU memory to make DP beneficial.

This PR (minor regression in short context but seems worth it given the saved memory boosts long-context and enables DP-attention, also the short context measurements are a bit noisy)

  backend  input_tokens  output_tokens  output_toks/s     req/s  median_itl_ms  median_ttft_ms
3    vllm          1000           1000    1323.397915  1.323398      29.755860     2307.676603
2    vllm          5000           1000    1041.455043  1.041455      33.205620     5491.457423
4    vllm         10000           1000     874.563079  0.874563      36.871404     8508.498527
1    vllm         32000           1000     190.195698  0.190196      35.948243   108055.433401

Baseline (#14769)

  backend  input_tokens  output_tokens  output_toks/s     req/s  median_itl_ms  median_ttft_ms
3    vllm          1000           1000    1380.973383  1.380973      30.029079     2223.776098
2    vllm          5000           1000    1047.717832  1.047718      33.303024     5499.557093
4    vllm         10000           1000     586.460476  0.586460      36.855329     8512.637936
1    vllm         32000           1000     162.816157  0.162816      42.981104   115515.806204

correctness tests:

VLLM_USE_V1=1  vllm serve /home/vllm-dev/DeepSeek-R1 --tensor-parallel-size 8 --trust-remote-code --disable-log-requests

lm_eval --model local-completions --tasks gsm8k --model_args model=/home/vllm-dev/DeepSeek-R1,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=5,max_retries=3,tokenized_request
s=False --limit 100

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

VLLM_USE_V1=0 vllm serve /home/vllm-dev/DeepSeek-R1 --tensor-parallel-size 8 --trust-remote-code --disable-log-requests

lm_eval --model local-completions --tasks gsm8k --model_args model=/home/vllm-dev/DeepSeek-R1,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=5,max_retries=3,tokenized_request
s=False --limit 100

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

hello，may I know what the parallel configuration of the baseline and ur test is like. I am also exploring the comparison before and after enabling matrix absorption preconditioning.