[ROCM] DSfp4 mla projection gemms weight dynamic quantization

[maleksan85]

Commands:

Server:

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FP4BMM=1 VLLM_ROCM_USE_AITER_MHA=0 VLLM_ROCM_USE_AITER_MLA=1 AMDGCN_USE_BUFFER_OPS=1 SAFETENSORS_FAST_GPU=1 vllm serve /data/models/DeepSeek-R1-0528-MXFP4-Preview --host localhost --port 8000 --tensor-parallel-size 8 --max-num-batched-tokens 32768 --trust-remote-code --no-enable-prefix-caching --disable-log-requests --gpu_memory_utilization 0.8 --async-scheduling --block-size 16 --load-format fastsafetensors --seed 123 --enforce-eager

Client:

curl http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
"model": "/data/models/DeepSeek-R1-0528-MXFP4-Preview",
"prompt": "San Francisco is a",
"max_tokens": 16,
"temperature": 0
}'

Output:

{"id":"cmpl-9afa69e020697740","object":"text_completion","created":1768286231,"model":"/data/models/DeepSeek-R1-0528-MXFP4-Preview","choices":[{"index":0,"text":" city known for its vibrant culture, stunning architecture, and diverse neighborhoods. Among its","logprobs":null,"finish_reason":"length","stop_reason":null,"token_ids":null,"prompt_logprobs":null,"prompt_token_ids":null}],"service_tier":null,"system_fingerprint":null,"usage":{"prompt_tokens":5,"total_tokens":21,"completion_tokens":16,"prompt_tokens_details":null},"kv_transfer_params":null}

legacy: VLLM_ROCM_USE_AITER_FP4BMM=0 VLLM_ROCM_USE_AITER_FP8BMM=0 also works

[gemini-code-assist]

Code Review

This pull request introduces support for DSfp4 MLA projection GEMMs with dynamic weight quantization on ROCm. The changes look good overall, adding new environment variables and logic to handle FP4 batched matrix multiplication. I've found a critical issue and a high-severity typo that should be addressed.

[mergify]

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

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

[cursor]

FP8 KV cache view skipped when FP4BMM enabled

High Severity

The condition fp8_attention and not self.is_aiter_triton_fp4_bmm_enabled can cause data corruption when both FP8 KV cache and FP4BMM are enabled simultaneously. The concat_and_cache_mla call at line 2068 writes FP8-encoded data when kv_cache_dtype is "fp8", but line 2077 skips the .view(fp8_dtype) conversion when is_aiter_triton_fp4_bmm_enabled is True. Since these conditions are independent (one depends on KV cache dtype, the other on weight dtype and feature flag), both can be True, causing subsequent operations to read FP8 bytes as the original dtype, producing garbage results.

 

[cursor]

Base class missing FP4BMM attribute setup in process_weights_after_loading

Medium Severity

MLACommonBaseImpl._v_up_proj was modified to use self.W_V and self.W_V_scale when is_aiter_triton_fp4_bmm_enabled is True, but MLACommonBaseImpl.process_weights_after_loading only sets these attributes for FP8BMM (in the if self.is_aiter_triton_fp8_bmm_enabled branch). When FP4BMM is enabled but FP8BMM is disabled, process_weights_after_loading falls into the else branch which sets W_UV and W_UK_T instead, leaving W_V and W_V_scale undefined. This causes the base class to be in an inconsistent state. While MLACommonImpl handles FP4BMM setup correctly, any direct subclass of MLACommonBaseImpl would hit an AttributeError at runtime.

Additional Locations (1)

• vllm/v1/attention/backends/mla/common.py#L1241-L1287