[Perf][vLLM-ATOM] Optimize Sparse MLA in vLLM-ATOM

[kliuae]

Motivation

This PR adds optimization for sparse MLA in vLLM-ATOM.

Technical Details

• Enable preshuffled indexer cache and add block_size 64
• Use persistent mla and add respective workspace buffers

Test Plan

Accuracy test with lm_eval on MI355X

Model: deepseek-ai/DeepSeek-V3.2

Server command:

ATOM_DISABLE_VLLM_PLUGIN=0 \
ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=0 \
VLLM_ROCM_USE_AITER=1 \
vllm serve deepseek-ai/DeepSeek-V3.2 \
  -tp 8 \
  --gpu-memory-utilization 0.8 \
  --no-enable-prefix-caching \
  --disable-uvicorn-access-log \
  --trust-remote-code \
  --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
  --kv-cache-dtype fp8

lm_eval command

lm_eval --model local-completions   --model_args model=deepseek-ai/DeepSeek-V3.2,base_url=http://localhost:8000/v1/completions,num_concurrent=64,tokenized_requests=False  --tasks gsm8k --num_fewshot 20

Benchmark command

vllm bench serve --model deepseek-ai/DeepSeek-V3.2 --port 8000 --dataset-name random --random-input-len 1024 --random-output-len 1024 --max-concurrency $CONC --num-prompts $(( CONC * 10 )) --ignore-eos --temperature 0

Test Result

lm_eval

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	_	0.9500	_	0.0060
		strict-match	20	exact_match	_	0.9484	_	0.0061

Performance on MI355X, TP8

ISL/OSL	Concurrency	KV Cache	PR req/s	Main req/s	PR over Main (req/s)	PR Total tok/s	Main Total tok/s	PR over Main (tok/s)
1k/1k	64	fp8	2.50	2.17	+15.21%	5123.55	4440.94	+15.37%
1k/1k	128	fp8	4.06	3.47	+17.00%	8320.14	7103.16	+17.13%

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[zejunchen-zejun]

Thank you @kliuae

Hi, @XiaobingSuper
Could you help review this PR?