[NVIDIA] Add Cutlass MLA backend

[kaixih]

This PR introduces the CUTLASS_MLA_VLLM_V1 backend, enabling support for ops.cutlass_mla_decode() on NVIDIA Blackwell GPUs.

It also includes performance results using DeepSeek-V3 on 8×B200 GPUs under DP+EP parallelism settings, which delivers ~17% improved throughput.

# With default triton backend:
============ Serving Benchmark Result ============
Successful requests:                     2989
Benchmark duration (s):                  1046.01
Total input tokens:                      2989000
Total generated tokens:                  2989000
Request throughput (req/s):              2.86
Output token throughput (tok/s):         2857.52
Total Token throughput (tok/s):          5715.04
---------------Time to First Token----------------
Mean TTFT (ms):                          200716.51
Median TTFT (ms):                        199463.35
P99 TTFT (ms):                           395239.25
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          826.04
Median TPOT (ms):                        826.20
P99 TPOT (ms):                           1001.39
---------------Inter-token Latency----------------
Mean ITL (ms):                           826.04
Median ITL (ms):                         648.89
P99 ITL (ms):                            8337.69
==================================================

With cutlass_mla backend:
============ Serving Benchmark Result ============
Successful requests:                     2989
Benchmark duration (s):                  881.52
Total input tokens:                      2989000
Total generated tokens:                  2989000
Request throughput (req/s):              3.39
Output token throughput (tok/s):         3390.73
Total Token throughput (tok/s):          6781.46
---------------Time to First Token----------------
Mean TTFT (ms):                          190244.11
Median TTFT (ms):                        189563.96
P99 TTFT (ms):                           372713.07
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          685.60
Median TPOT (ms):                        686.96
P99 TPOT (ms):                           858.01
---------------Inter-token Latency----------------
Mean ITL (ms):                           685.60
Median ITL (ms):                         518.56
P99 ITL (ms):                            7738.23
==================================================

To repro the results:

# Server side with triton backend (Plz use VLLM_ATTENTION_BACKEND=CUTLASS_MLA_VLLM_V1 for cutlass backend):
VLLM_LOGGING_LEVEL=DEBUG \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
  vllm serve deepseek-ai/DeepSeek-V3 \
    --trust-remote-code \
    --max-model-len=2048 \
    --block-size=128 \
    --max-num-seqs=512 \
    --gpu_memory_utilization=0.97 \
    --data-parallel-size $NUM_GPUS --enable-expert-parallel \
    --disable-log-requests

# client side:
python $VLLM_PATH/benchmarks/benchmark_serving.py \
  --model deepseek-ai/DeepSeek-V3 \
  --dataset-name random \
  --ignore-eos \
  --num-prompts 3000 \
  --max-concurrency 3000 \
  --random-input-len 1000 \
  --random-output-len 1000

cc. @kushanam

[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.

🚀

[mgoin]

Make sure to update the messages for CutlassMLA

[kaixih]

Done. PTAL.

[LucasWilkinson]

The perf is looking really good! thanks for the contribution!

Do you mind doing accuracy checks

VLLM_ATTENTION_BACKEND=CUTLASS_MLA_VLLM_V1 VLLM_USE_V1=1 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 --task gsm8k --num_fewshot 5  --batch_size auto

[LucasWilkinson]

we should leave

self.aot_schedule = is_vllm_fa and (get_flash_attn_version() == 3)

and updated:

if current_platform.is_cuda():
     self.page_size = self.runner.block_size

to

if self.aot_schedule:
     self.page_size = self.runner.block_size

[kaixih]

Do you mean we should do:

self.aot_schedule = is_vllm_fa and (get_flash_attn_version() == 3)
if current_platform.is_cuda():
     self.page_size = self.runner.block_size

(... since the self.aot_schedule might be used elsewhere?)

[LucasWilkinson]

oh sorry ya aot_schedule is not actually used, aot_schedule stands for ahead-of-time schedule and was meant to indicate if we should use the ahead of time scheduler for flash-attn (something I may add, when I updated FA I had added it for the FA (

vllm/vllm/v1/attention/backends/flash_attn.py

Lines 341 to 354 in 2f925e5

	if self.aot_schedule:
	return get_scheduler_metadata(
	batch_size=batch_size,
	max_seqlen_q=max_query_len,
	max_seqlen_k=max_seq_len,
	cache_seqlens=seqlens,
	num_heads_q=self.num_heads_q,
	num_heads_kv=self.num_heads_kv,
	headdim=self.headdim,
	page_size=self.page_size,
	cu_seqlens_q=cu_query_lens,
	causal=causal,
	window_size=self.aot_sliding_window,
	)

) and MLA backends but ripped it out of MLA before the PR landed since there was issues, you can just remove self.aot_schedule = is_vllm_fa and (get_flash_attn_version() == 3) completely. Good catch!

[LucasWilkinson]

the perf looks really good! I think we should turn this on by default for blackwell

[LucasWilkinson]

does this need to be torch.zeros? or does torch.empty work (faster)

[kaixih]

The perf is looking really good! thanks for the contribution!

Do you mind doing accuracy checks

VLLM_ATTENTION_BACKEND=CUTLASS_MLA_VLLM_V1 VLLM_USE_V1=1 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 --task gsm8k --num_fewshot 5  --batch_size auto

Do you know any commend to test a model with num_heads = 128? And probably no TP.

[LucasWilkinson]

Do you know any commend to test a model with num_heads = 128? And probably no TP.

Not that im aware of :/ this is the smallest MLA model I am aware of

[mgoin]

This is the smallest model with MLA ( @tlrmchlsmth found it the other day) https://huggingface.co/deepseek-ai/deepseek-vl2-tiny

[tlrmchlsmth]

This is the smallest model with MLA ( @tlrmchlsmth found it the other day) https://huggingface.co/deepseek-ai/deepseek-vl2-tiny

just want to clarify that the way I found it was grep -r deepseek tests | grep tiny

[LucasWilkinson]

Ah I don't think it's an MLA model :/

    "kv_lora_rank": null,
    ...
    "use_mla": false,

[LucasWilkinson]

LGTM, other than I do think we should turn it on by default for Blackwell: https://github.com/vllm-project/vllm/pull/17625/files#r2075575324

Any reason not to?

Edit: oh and ideally id still like to see accuracy numbers, i.e.:

VLLM_ATTENTION_BACKEND=CUTLASS_MLA_VLLM_V1 VLLM_USE_V1=1 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 --task gsm8k --num_fewshot 5  --batch_size auto

[kaixih]

Edit: oh and ideally id still like to see accuracy numbers...

@LucasWilkinson this DeepSeek-V2-Lite-Chat only has attention head number == 16 and --tp=2 is not ok. Any advice on a model using head num == 128?

[kaixih]

other than I do think we should turn it on by default for Blackwell, Any reason not to?

My main concern is that the CUTLASS MLA kernel has more limited support compared to the Triton version. For example, num_heads must be 128, and the page table must be padded to a multiple of 128 / page_size. We're working on expanding the supported use cases, but until then, we'd prefer to keep this as an optional backend. What do you think?

[LucasWilkinson]

other than I do think we should turn it on by default for Blackwell, Any reason not to?

My main concern is that the CUTLASS MLA kernel has more limited support compared to the Triton version. For example, num_heads must be 128, and the page table must be padded to a multiple of 128 / page_size. We're working on expanding the supported use cases, but until then, we'd prefer to keep this as an optional backend. What do you think?

Oh interesting, can we turn it on my default for models with head dim 128? as for block size I think we can force it similar to to FlashMLA (

vllm/vllm/platforms/cuda.py

Lines 145 to 158 in 5e6f939

	# TODO(lucas): handle this more gracefully
	# Note: model_config may be None during testing
	if model_config is not None and model_config.use_mla:
	# if `VLLM_ATTENTION_BACKEND` is not set and we are using MLA, then
	# we default to FlashMLA backend, so we need to force the blocksize
	# here
	use_flashmla = (envs.VLLM_ATTENTION_BACKEND is None \
	or envs.VLLM_ATTENTION_BACKEND == "FLASHMLA")
	from vllm.attention.ops.flashmla import is_flashmla_supported
	if use_flashmla and is_flashmla_supported()[0] \
	and cache_config.block_size != 64:
	cache_config.block_size = 64
	logger.info(
	"Forcing kv cache block size to 64 for FlashMLA backend.")

), we plan to eventually handle this more gracefully (i.e. better support of the attention backend being able to specify the block size it wants if the user has not specified one)

Im ok landing this for now and opening a new PR to turn it on by default; but I think we still should, we should strive to make the best they can be.

[LucasWilkinson]

Edit: oh and ideally id still like to see accuracy numbers...

@LucasWilkinson this DeepSeek-V2-Lite-Chat only has attention head number == 16 and --tp=2 is not ok. Any advice on a model using head num == 128?

Ah the only one I know of is DeepSeek V3 or R1, if you have a system that can run that then those accuracy numbers would be good to see!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kaixih.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[kaixih]

Got the lm_eval output with the cutlass backend, which matches the triton backend:

|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|

How to repro:

# server side: bash run.sh [cutlass]
if [[ "$1" == "cutlass" ]]; then
  export VLLM_ATTENTION_BACKEND=CUTLASS_MLA_VLLM_V1
fi

VLLM_LOGGING_LEVEL=DEBUG \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
  vllm serve deepseek-ai/DeepSeek-V3 \
    --trust-remote-code \
    --max-model-len=4096 \
    --block-size=128 \
    --max-num-seqs=512 \
    --gpu_memory_utilization=0.97 \
    --data-parallel-size $NUM_GPUS --enable-expert-parallel \
    --disable-log-requests &

# client side
lm_eval --model local-completions --tasks gsm8k \
    --model_args model=deepseek-ai/DeepSeek-V3,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=5,max_retries=3,tokenized_requests=False \
    --limit 100

[tlrmchlsmth]

Is this right? I thought we still needed to guard against using aot_schedule if using FA2

[kaixih]

See comment here. My understanding is that this aot_schedule is irrelevant here. @LucasWilkinson can correct me.

[LucasWilkinson]

Ya this is right, page-size is used by chunked-prefill in MLA, this naming was a bug

[kaixih]

Do I need to make any further changes, or is this good to go now?

[tlrmchlsmth]

This PR is pretty red... these tests aren't all failing on main are they?

[kaixih]

This PR is pretty red... these tests aren't all failing on main are they?

I am checking the logs but nothing seems to be related to my change. Can you please advise? @tlrmchlsmth

[kaixih]

@tlrmchlsmth ok. I think i have done all I can do. Basically I have reverted all my changes that are for the existing files (only leave one that adding a new file). And you can see the tests are still red. Can you help?

[kaixih]

Force push the original PR.

[kaixih]

@tlrmchlsmth Can you please advise? ^^

[houseroad]

I think we have fixed the CI, could you rebase the PR again?

[kaixih]

@houseroad Thanks, I've just rebased. However, the Lint and Deploy Charts / lint-and-deploy (pull_request) check is failing, and it appears to be unrelated to the changes in this PR. Can you advise?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kaixih.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork