Skip to content

[ROCm][Deepseek] dsv3.2 further optimization #41217

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tjtanaa merged 5 commits into from
May 1, 2026
Merged

[ROCm][Deepseek] dsv3.2 further optimization #41217

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
36bd605
further optimize dsv3.2
ganyi1996ppo Nov 27, 2025
4df6432
add some assertion for non-supported case
ganyi1996ppo Apr 29, 2026
999ca3a
fix precommit
ganyi1996ppo Apr 30, 2026
3a6d3f5
Mark FP8 KV dtype as supported in ROCm AITER Sparse MLA
mawong-amd Apr 30, 2026
f77e120
resolve comments and support block-size 1
ganyi1996ppo May 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/design/attention_backends.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ configuration.
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
61 changes: 38 additions & 23 deletions vllm/model_executor/models/deepseek_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -674,30 +674,45 @@ def forward(
) -> torch.Tensor:
q, _ = self.wq_b(qr)
q = q.view(-1, self.n_head, self.head_dim)
q_pe, q_nope = torch.split(
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
)
# Fused wk + weights_proj: one GEMM, then split
kw, _ = self.wk_weights_proj(hidden_states)
k = kw[:, : self.head_dim]
weights = kw[:, self.head_dim :]

k = self.k_norm(k)
k_pe, k_nope = torch.split(
k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
)

q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
# Note: RoPE (NeoX) can introduce extra leading dimensions during compilation
# so we need to reshape back to token-flattened shapes
q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
k_pe = k_pe.reshape(-1, 1, self.rope_dim)

# `rotary_emb` is shape-preserving; `q_pe` is already
# [num_tokens, n_head, rope_dim].
q = torch.cat([q_pe, q_nope], dim=-1)
# `k_pe` is [num_tokens, 1, rope_dim] (MQA).
k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
if current_platform.is_rocm():
# This path should works on all platform, will remove extra
# branches in the future
# Fused wk + weights_proj: one GEMM, then split
kw, _ = self.wk_weights_proj(hidden_states)
k = kw[:, : self.head_dim]
weights = kw[:, self.head_dim :]

k = self.k_norm(k)

rotary_emb(
positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1)
)
else:
q_pe, q_nope = torch.split(
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
)
# Fused wk + weights_proj: one GEMM, then split
kw, _ = self.wk_weights_proj(hidden_states)
k = kw[:, : self.head_dim]
weights = kw[:, self.head_dim :]

k = self.k_norm(k)
k_pe, k_nope = torch.split(
k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
)

q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
# Note: RoPE (NeoX) can introduce extra leading dimensions during
# compilation so we need to reshape back to token-flattened shapes
q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
k_pe = k_pe.reshape(-1, 1, self.rope_dim)

# `rotary_emb` is shape-preserving; `q_pe` is already
# [num_tokens, n_head, rope_dim].
q = torch.cat([q_pe, q_nope], dim=-1)
# `k_pe` is [num_tokens, 1, rope_dim] (MQA).
k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)

# we only quant q here since k quant is fused with cache insertion
q = q.view(-1, self.head_dim)
Expand Down
2 changes: 1 addition & 1 deletion vllm/v1/attention/backends/mla/indexer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ def get_name() -> str:

@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [1 if current_platform.is_rocm() else 64]
return [1, 64] if current_platform.is_rocm() else [64]

@classmethod
def get_supported_head_sizes(cls) -> list[int]:
Expand Down
4 changes: 4 additions & 0 deletions vllm/v1/attention/backends/mla/rocm_aiter_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,7 @@ class AiterMLAHelper:
"""

_AITER_MIN_MLA_HEADS: Final = 16
_AITER_UNSUPPORTED_HEADS = [32]
Copy link
Copy Markdown
Collaborator

@tjtanaa tjtanaa Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it just head size 32 that has issue?

So can I understand as

head size 16 is fine; head size 32 is NOT supported;
head size 64 is fine; head size 128 is fine as well, etc .

Copy link
Copy Markdown
Contributor Author

@ganyi1996ppo ganyi1996ppo Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like so, I found when head_size equals to 32, the aiter will automatically pick the, as the symbol name suggested, 16head implementation and cause the illegal memory access. I'm not sure if other shape have the same issue, I only tested tp8 and tp4 case


@staticmethod
def check_num_heads_validity(num_heads: int):
Expand All @@ -419,6 +420,9 @@ def get_actual_mla_num_heads(num_heads: int) -> int:

@staticmethod
def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
assert num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS, (
f"unsupported head_num: {num_heads}"
)
Comment on lines +423 to +425
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The assertion assert num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS is added to get_mla_padded_q, but this check should ideally be performed at the entry point of the attention layer or during configuration validation to prevent unnecessary runtime checks in the hot path of the forward pass.

return (
q
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
Expand Down
Loading
Loading