minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Separate ROCm-specific DeepSeek MHA and MLA forward paths (#31531)

Browse Source
This commit is contained in:
Diya Peng
2026-08-08 13:39:36 -07:00
committed by GitHub
parent 3fbb5330c7
commit c4f018ba1d
12 changed files with 1366 additions and 590 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/operations.py
+1 -1
View File
@@ -75,7 +75,7 @@ def _resolve_tbo_child_contexts():
"""Return (child_ctx_a, child_ctx_b) derived from the active TboAttnBackend,
or (None, None) if the active backend is not a TBO dispatcher (e.g. a
backend that handles TBO splitting internally like DeepSeek MHA's
_resolve_attn_backend path)."""
resolve_attn_backend path)."""
# Lazy import to avoid circular dependency at module load time.
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
python/sglang/srt/models/deepseek_common/attention_backend_handler.py
+17
View File
@@ -16,6 +16,23 @@ from sglang.srt.utils import is_sm100_or_sm110_supported, use_intel_amx_backend
MHA_ONE_SHOT_SUPPORTED_BACKENDS = ["fa3", "flashinfer", "flashmla"]
# ROCm runs dedicated MHA/MLA implementations (forward_mha_rocm.py /
# forward_mla_rocm.py) so the shared CUDA paths carry no AMD branches. Backend
# handlers keep returning the generic method; the platform swap happens here.
# MHA_CHUNKED_KV has no ROCm entry because its accumulation step needs the
# CUDA-only merge_state_v2 kernel.
_ROCM_FORWARD_METHODS = {
AttnForwardMethod.MHA: AttnForwardMethod.MHA_ROCM,
AttnForwardMethod.MHA_ONE_SHOT: AttnForwardMethod.MHA_ONE_SHOT_ROCM,
AttnForwardMethod.MLA: AttnForwardMethod.MLA_ROCM,
}
def resolve_rocm_forward_method(method: AttnForwardMethod) -> AttnForwardMethod:
if not _is_hip:
return method
return _ROCM_FORWARD_METHODS.get(method, method)
class AttentionBackendRegistry:
_handlers = {}
python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py
+5 -1
View File
@@ -1,13 +1,17 @@
from .forward_methods import AttnForwardMethod
from .forward_mha import DeepseekMHAForwardMixin
from .forward_mha_rocm import DeepseekMHARocmForwardMixin
from .forward_mla import DeepseekMLAForwardMixin
from .forward_mla_fused_rope_cpu import DeepseekMLACpuForwardMixin
from .forward_mla_fused_rope_rocm import DeepseekMLARocmForwardMixin
from .forward_mla_fused_rope_rocm import DeepseekMLAFusedRopeRocmForwardMixin
from .forward_mla_rocm import DeepseekMLARocmForwardMixin
__all__ = [
"AttnForwardMethod",
"DeepseekMHAForwardMixin",
"DeepseekMHARocmForwardMixin",
"DeepseekMLACpuForwardMixin",
"DeepseekMLAForwardMixin",
"DeepseekMLAFusedRopeRocmForwardMixin",
"DeepseekMLARocmForwardMixin",
]
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py
+9
View File
@@ -30,3 +30,12 @@ class AttnForwardMethod(IntEnum):
# Use Deepseek V3.2 sparse multi-latent attention for NPU
DSA_NPU = auto()
# Use multi-head attention for ROCm
MHA_ROCM = auto()
# Use one-shot multi-head attention for ROCm
MHA_ONE_SHOT_ROCM = auto()
# Use absorbed multi-latent attention for ROCm
MLA_ROCM = auto()
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+18 -129
View File
@@ -14,9 +14,6 @@ from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mha_extend,
filter_dcp_local_kv_indices,
)
from sglang.srt.layers.quantization.fp8_utils import (
materialize_bpreshuffle_fp8_scale_tuple,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import (
get_attn_backend,
@@ -24,22 +21,15 @@ from sglang.srt.model_executor.forward_context import (
)
from sglang.srt.models.deepseek_common.utils import (
_is_cuda,
_is_hip,
_is_musa,
_is_npu,
_use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95,
)
from sglang.srt.runtime_context import (
get_exec,
get_parallel,
get_schedule,
)
from sglang.srt.utils import BumpAllocator, get_bool_env_var, next_power_of_2
_use_fp8_prefill_attn = (
get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and _use_aiter_gfx95
)
from sglang.srt.utils import BumpAllocator, next_power_of_2
if TYPE_CHECKING:
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
@@ -51,21 +41,15 @@ if _is_cuda:
elif _is_musa:
from sgl_kernel import concat_mla_k
if _use_aiter_gfx95:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype
from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
def _resolve_attn_backend(forward_batch: ForwardBatch):
def resolve_attn_backend(forward_batch: ForwardBatch):
backend = get_attn_backend()
if isinstance(backend, TboAttnBackend):
backend = backend.primary
return backend
def _forward_dsa_indexer_for_mha(
def forward_dsa_indexer_for_mha(
indexer,
*,
hidden_states: torch.Tensor,
@@ -173,39 +157,12 @@ class DeepseekMHAForwardMixin:
# DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
if self.use_dsa:
# DSA requires unquantized q_lora for the indexer. When q_b_proj is FP8
# on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
# the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
# tuple.
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q_quanted = materialize_bpreshuffle_fp8_scale_tuple(q_quanted)
q = self.q_b_proj(q_quanted)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
else:
q_lora = self.q_a_layernorm(q)
q = self.q_b_proj(q_lora)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
q_lora = self.q_a_layernorm(q)
q = self.q_b_proj(q_lora)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
if self.should_run_indexer():
_forward_dsa_indexer_for_mha(
forward_dsa_indexer_for_mha(
self.indexer,
hidden_states=hidden_states,
q_lora=q_lora,
@@ -213,34 +170,6 @@ class DeepseekMHAForwardMixin:
forward_batch=forward_batch,
layer_id=self.layer_id,
)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
# MXFP4: fused RMSNorm + quant
q, _, _, _ = fused_rms_mxfp4_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
q, _, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q = materialize_bpreshuffle_fp8_scale_tuple(q)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
else:
q = self.q_a_layernorm(q)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
@@ -255,25 +184,7 @@ class DeepseekMHAForwardMixin:
kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
latent_cache = latent_cache.unsqueeze(1)
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
kv_a,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True, # return unqaunt kv_a
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
kv_a_quanted = materialize_bpreshuffle_fp8_scale_tuple(kv_a_quanted)
else:
kv_a = self.kv_a_layernorm(kv_a)
kv_a = self.kv_a_layernorm(kv_a)
k_pe = latent_cache[:, :, self.kv_lora_rank :]
@@ -281,7 +192,7 @@ class DeepseekMHAForwardMixin:
# (fused RoPE + quantize for Q/K, direct FP8 KV-cache write) and
# returns FP8 tensors ready for its kernel. Backends without the
# hook fall through to the BF16 path below.
backend = _resolve_attn_backend(forward_batch)
backend = resolve_attn_backend(forward_batch)
if hasattr(backend, "prepare_prefill_qkv"):
q_out, k_out, v_out = backend.prepare_prefill_qkv(
q=q,
@@ -333,31 +244,12 @@ class DeepseekMHAForwardMixin:
q.dtype,
forward_batch,
)
if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8:
# MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
# into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
# directly in FP8, avoiding a separate elementwise cast
k, v = self.kv_b_proj(
(
kv_a,
k_pe.expand(-1, self.num_local_heads, -1),
self.qk_nope_head_dim,
self.v_head_dim,
fp8_dtype,
)
)[0]
else:
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
kv = self.kv_b_proj(kv_a_quanted)[0]
else:
kv = self.kv_b_proj(kv_a)[0]
kv = kv.view(
-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
)
k_nope = kv[..., : self.qk_nope_head_dim]
v = kv[..., self.qk_nope_head_dim :]
kv = self.kv_b_proj(kv_a)[0]
kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
k_nope = kv[..., : self.qk_nope_head_dim]
v = kv[..., self.qk_nope_head_dim :]
k = self._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
k = self._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
return q, k, v, forward_batch
def forward_normal_core(
@@ -464,7 +356,7 @@ class DeepseekMHAForwardMixin:
forward_batch: ForwardBatch,
) -> torch.Tensor:
# kv_b_proj needs BF16 input, but legacy q.dtype was BF16 by accident.
backend = _resolve_attn_backend(forward_batch)
backend = resolve_attn_backend(forward_batch)
pack_fn = getattr(backend, "pack_prefix_chunk_kv", None)
kv_a_dtype = torch.bfloat16 if pack_fn is not None else q.dtype
@@ -531,7 +423,7 @@ class DeepseekMHAForwardMixin:
k_pe: torch.Tensor,
forward_batch: ForwardBatch,
):
if _is_cuda or _use_aiter_gfx95:
if _is_cuda:
# Save latent cache
get_token_to_kv_pool().set_mla_kv_buffer(
self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
@@ -556,7 +448,7 @@ class DeepseekMHAForwardMixin:
dst_dtype: torch.dtype,
forward_batch: ForwardBatch,
):
if _is_cuda or _use_aiter_gfx95:
if _is_cuda:
kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
self.attn_mha, kv_indices, dst_dtype
@@ -632,9 +524,6 @@ class DeepseekMHAForwardMixin:
attn_dtype = k_nope.dtype
k = k_nope.new_empty(*k_shape, dtype=attn_dtype)
concat_and_cast_mha_k_triton(k, k_nope, k_pe)
elif _is_hip and self.current_attention_backend == "aiter":
k = k_nope.new_empty(*k_shape)
concat_and_cast_mha_k_triton(k, k_nope, k_pe)
else:
k = k_nope.new_empty(*k_shape)
k[..., : self.qk_nope_head_dim] = k_nope
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py
+324
View File
@@ -0,0 +1,324 @@
"""AMD/ROCm multi-head attention forward path for DeepSeek models.
`AttnForwardMethod.MHA_ROCM` and `AttnForwardMethod.MHA_ONE_SHOT_ROCM` route
here, which keeps every AITER/gfx95 kernel choice out of the shared
`forward_mha.py`. Only the prepare step is platform specific; the attention
cores in `DeepseekMHAForwardMixin` are reused as-is.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.kernels.ops.attention.utils import concat_and_cast_mha_k_triton
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mha_extend,
filter_dcp_local_kv_indices,
)
from sglang.srt.layers.quantization.fp8_utils import (
materialize_bpreshuffle_fp8_scale_tuple,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import (
forward_dsa_indexer_for_mha,
resolve_attn_backend,
)
from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95,
)
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.utils import BumpAllocator, get_bool_env_var
if TYPE_CHECKING:
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
_use_fp8_prefill_attn = (
get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and _use_aiter_gfx95
)
if _use_aiter_gfx95:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype
from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
class DeepseekMHARocmForwardMixin:
def forward_normal_rocm_prepare(
self: DeepseekV2AttentionMLA,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
):
if self.q_lora_rank is not None:
q, latent_cache = (
get_attn_tp_context()
.fetch_qkv_latent()
.split(
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
dim=-1,
)
)
# DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
if self.use_dsa:
# DSA requires unquantized q_lora for the indexer. When q_b_proj is FP8
# on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
# the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
# tuple.
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q_quanted = materialize_bpreshuffle_fp8_scale_tuple(q_quanted)
q = self.q_b_proj(q_quanted)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
else:
q_lora = self.q_a_layernorm(q)
q = self.q_b_proj(q_lora)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
if self.should_run_indexer():
forward_dsa_indexer_for_mha(
self.indexer,
hidden_states=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
# MXFP4: fused RMSNorm + quant
q, _, _, _ = fused_rms_mxfp4_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
q, _, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q = materialize_bpreshuffle_fp8_scale_tuple(q)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
else:
q = self.q_a_layernorm(q)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
else:
q = self.q_proj(hidden_states)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
_, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
latent_cache = latent_cache.unsqueeze(1)
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
kv_a,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True, # return unqaunt kv_a
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
kv_a_quanted = materialize_bpreshuffle_fp8_scale_tuple(kv_a_quanted)
else:
kv_a = self.kv_a_layernorm(kv_a)
k_pe = latent_cache[:, :, self.kv_lora_rank :]
# Backend prefill hook: the backend owns the BF16->FP8 transition
# (fused RoPE + quantize for Q/K, direct FP8 KV-cache write) and
# returns FP8 tensors ready for its kernel. Backends without the
# hook fall through to the BF16 path below.
backend = resolve_attn_backend(forward_batch)
if hasattr(backend, "prepare_prefill_qkv"):
q_out, k_out, v_out = backend.prepare_prefill_qkv(
q=q,
q_pe=q_pe,
kv_a=kv_a,
k_pe=k_pe,
positions=positions,
layer=self,
forward_batch=forward_batch,
)
return q_out, k_out, v_out, forward_batch
if self.rotary_emb is not None:
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
q[..., self.qk_nope_head_dim :] = q_pe
self._set_mla_kv_buffer_rocm(latent_cache, kv_a, k_pe, forward_batch)
if (
forward_batch.mha_one_shot
and sum(forward_batch.extend_prefix_lens_cpu) != 0
):
if (
self.use_dsa
and self.kv_cache_dtype == "fp8_e4m3"
and (
not get_exec().kernel.dsa_decode_backend == "trtllm"
or not get_exec().kernel.dsa_prefill_backend == "trtllm"
)
):
# FP8 path: dequantize DSA-specific FP8 format to BF16
kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch)
else:
# BF16/FP16 path: directly fetch from cache
if get_parallel().dcp_enabled:
kv_a, k_pe = all_gather_kv_cache_for_mha_extend(
get_token_to_kv_pool(),
self.attn_mha,
forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices,
forward_batch.seq_lens,
forward_batch.extend_prefix_lens,
forward_batch.extend_prefix_lens_cpu,
forward_batch.extend_seq_lens,
kv_a,
k_pe,
)
else:
kv_a, k_pe = self._get_mla_kv_buffer_rocm(
forward_batch.fetch_mha_one_shot_kv_indices(),
q.dtype,
forward_batch,
)
if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8:
# MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
# into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
# directly in FP8, avoiding a separate elementwise cast
k, v = self.kv_b_proj(
(
kv_a,
k_pe.expand(-1, self.num_local_heads, -1),
self.qk_nope_head_dim,
self.v_head_dim,
fp8_dtype,
)
)[0]
else:
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
kv = self.kv_b_proj(kv_a_quanted)[0]
else:
kv = self.kv_b_proj(kv_a)[0]
kv = kv.view(
-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
)
k_nope = kv[..., : self.qk_nope_head_dim]
v = kv[..., self.qk_nope_head_dim :]
k = self._concat_and_cast_mha_k_rocm(k_nope, k_pe)
return q, k, v, forward_batch
def forward_normal_one_shot_rocm_prepare(
self: DeepseekV2AttentionMLA,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
):
forward_batch.mha_one_shot = True
return self.forward_normal_rocm_prepare(
positions, hidden_states, forward_batch, zero_allocator
)
def _concat_and_cast_mha_k_rocm(
self: DeepseekV2AttentionMLA,
k_nope: torch.Tensor,
k_pe: torch.Tensor,
):
k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim)
k = k_nope.new_empty(*k_shape)
if self.current_attention_backend == "aiter":
concat_and_cast_mha_k_triton(k, k_nope, k_pe)
else:
k[..., : self.qk_nope_head_dim] = k_nope
k[..., self.qk_nope_head_dim :] = k_pe
return k
def _set_mla_kv_buffer_rocm(
self: DeepseekV2AttentionMLA,
latent_cache: torch.Tensor,
kv_a: torch.Tensor,
k_pe: torch.Tensor,
forward_batch: ForwardBatch,
):
if _use_aiter_gfx95:
get_token_to_kv_pool().set_mla_kv_buffer(
self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
)
else:
latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1)
latent_cache[:, :, self.kv_lora_rank :] = k_pe.clone()
get_token_to_kv_pool().set_kv_buffer(
self.attn_mha, forward_batch.out_cache_loc, latent_cache, None
)
def _get_mla_kv_buffer_rocm(
self: DeepseekV2AttentionMLA,
kv_indices: torch.Tensor,
dst_dtype: torch.dtype,
forward_batch: ForwardBatch,
):
if _use_aiter_gfx95:
kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
self.attn_mha, kv_indices, dst_dtype
)
kv_a = kv_a.squeeze(1)
else:
latent_cache_buf = get_token_to_kv_pool().get_key_buffer(
self.attn_mha.layer_id
)
latent_cache = latent_cache_buf[kv_indices].contiguous().to(dst_dtype)
kv_a, k_pe = latent_cache.split(
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
)
kv_a = kv_a.squeeze(1).contiguous()
return kv_a, k_pe
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+77 -454
View File
@@ -8,7 +8,6 @@ import torch
from sglang.kernels.ops.kvcache.cache_ops import absorbed_bmm_concat_cast_q_fp8
from sglang.kernels.ops.quantization.fp8_kernel import (
fp8_dtype,
per_tensor_quant_mla_fp8,
per_token_group_quant_mla_deep_gemm_masked_fp8,
)
@@ -28,9 +27,6 @@ from sglang.srt.layers.dcp import (
dcp_a2a_lse_reduce,
)
from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
from sglang.srt.layers.quantization.fp8_utils import (
materialize_bpreshuffle_fp8_scale_tuple,
)
from sglang.srt.layers.radix_attention import unified_attention_with_output
from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
from sglang.srt.lora.deepseek_mla_correction import (
@@ -62,12 +58,8 @@ from sglang.srt.models.deepseek_common.utils import (
_is_cpu,
_is_cublas_ge_129,
_is_cuda,
_is_gfx95_supported,
_is_hip,
_is_musa,
_use_aiter,
_use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95,
)
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.state_capturer.indexer_topk import (
@@ -101,7 +93,16 @@ def _select_local_dcp_heads_for_autotune(
return attn_output.narrow(1, rank * num_local_heads, num_local_heads)
def _is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
def is_dcp_mla_decode_phase(forward_batch: ForwardBatch) -> bool:
if not get_parallel().dcp_enabled:
return False
return (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
def is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
# FlashMLA exposes natural-log softmax LSE. FlashInfer MLA and the other
# currently supported MLA DCP decode backends expose base-2 LSE.
return attention_backend == "flashmla"
@@ -111,58 +112,7 @@ if _is_cuda:
from sglang.kernels.ops.gemm import bmm_fp8
if _use_aiter:
# aiter ROCm/aiter#2958 renamed the public `fused_qk_rmsnorm` in
# `aiter.ops.fused_qk_norm_rope_cache_quant` to a private `_fused_qk_rmsnorm`
# and introduced a unified entry point in `aiter.ops.fused_qk_rmsnorm_group_quant`
# with a different (in-place, kwarg-only, no-return) signature. Probe for the
# new symbol first so SGLang works with both pre- and post-#2958 aiter without
# requiring the docker pin to be bumped atomically.
try:
from aiter.ops.enum import QuantType as _AiterQuantType
from aiter.ops.fused_qk_rmsnorm_group_quant import (
fused_qk_rmsnorm as _aiter_fused_qk_rmsnorm_unified,
)
def fused_qk_rmsnorm_bf16(q, q_weight, q_eps, k, k_weight, k_eps):
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
_aiter_fused_qk_rmsnorm_unified(
q_out_quantized=q_out,
k_out=k_out,
q=q,
q_weight=q_weight,
q_epsilon=q_eps,
k=k,
k_weight=k_weight,
k_epsilon=k_eps,
quant_type=_AiterQuantType.No,
)
return q_out, k_out
except ImportError:
from aiter.ops.fused_qk_norm_rope_cache_quant import (
fused_qk_rmsnorm as fused_qk_rmsnorm_bf16,
)
from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
)
if _use_aiter_gfx95:
from aiter.ops.triton.fused_fp8_quant import (
fused_flatten_fp8_group_quant,
fused_rms_fp8_group_quant,
)
from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
batched_gemm_afp4wfp4_pre_quant,
fused_flatten_mxfp4_quant,
fused_rms_mxfp4_quant,
)
from sglang.srt.layers.rocm_linear_utils import fused_qk_rope_cat_and_cache_mla
def _should_defer_dsa_cp_kv_gather(
def should_defer_dsa_cp_kv_gather(
*,
dsa_prefill_cp: bool,
fuse_rope_for_trtllm_mla: bool,
@@ -207,7 +157,7 @@ class DeepseekMLAForwardMixin:
return False
if not self.use_dsa:
return False
if self.use_deep_gemm_bmm or _is_hip:
if self.use_deep_gemm_bmm:
return False
if is_kv_b_lora_active(self):
return False
@@ -290,12 +240,6 @@ class DeepseekMLAForwardMixin:
return None
if self._fuse_rope_for_trtllm_mla(forward_batch):
return None
if self._skip_rope_for_dsa_tilelang_fused():
return None
if self._skip_rope_for_aiter_fused_mla():
return None
if _use_aiter and _is_gfx95_supported and not self.use_dsa:
return None
# Graph/compile surfaces run their own dispatch; the python-side
# stash handshake is eager-only.
if is_graph_dsa_split_op_surface(forward_batch):
@@ -356,11 +300,7 @@ class DeepseekMLAForwardMixin:
# weights and skip the per-layer Q all-gather (bf16 decode absorb only).
q_replicate_active = (
get_parallel().dcp_replicate_q_proj
and get_parallel().dcp_enabled
and (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
and is_dcp_mla_decode_phase(forward_batch)
and not self.use_deep_gemm_bmm
and self.w_kc_qrep is not None
and self.q_b_proj_qrep_weight is not None
@@ -394,69 +334,8 @@ class DeepseekMLAForwardMixin:
k_nope = self.kv_a_layernorm(k_nope)
current_stream.wait_stream(self.alt_stream)
else:
if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
q, _, k_nope, *_ = fused_rms_mxfp4_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
)
else:
q_lora = None
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
if self.use_dsa:
q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q_quanted = materialize_bpreshuffle_fp8_scale_tuple(
q_quanted
)
q = q_quanted
else:
q, _, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q = materialize_bpreshuffle_fp8_scale_tuple(q)
elif _use_aiter:
q, k_nope = fused_qk_rmsnorm_bf16(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
# q_lora needed by indexer
if self.use_dsa:
@@ -639,50 +518,6 @@ class DeepseekMLAForwardMixin:
expected_m,
)
q_nope_out = q_nope_out[:, :expected_m, :]
elif _is_hip:
# TODO(haishaw): add bmm_fp8 to ROCm
if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8:
x = q_nope.transpose(0, 1)
q_nope_out = torch.empty(
x.shape[0],
x.shape[1],
self.w_kc.shape[2],
device=x.device,
dtype=torch.bfloat16,
)
batched_gemm_afp4wfp4_pre_quant(
x,
self.w_kc.transpose(-2, -1),
self.w_scale_k.transpose(-2, -1),
torch.bfloat16,
q_nope_out,
)
else:
if (
_use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn
) or (
get_is_capture_mode()
and self.w_kc.dtype == torch.float8_e4m3fnuz
):
# fp8 Triton kernel: always on gfx950,
# cudagraph-only on gfx942 (hides launch overhead)
q_nope_out = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
X=q_nope,
WQ=self.w_kc.transpose(-1, -2),
w_scale=self.w_scale,
group_size=128,
YQ=None, # allocate (B, M, N)
transpose_bm=False, # (B, M, N)
transpose_bm_in=True, # (M, B, K)
dtype=torch.bfloat16,
)
else:
q_nope_out = torch.bmm(
q_nope.to(torch.bfloat16).transpose(0, 1),
self.w_kc.to(torch.bfloat16) * self.w_scale,
)
elif self.w_kc.dtype == torch.float8_e4m3fn:
if _is_cpu:
q_nope_out = torch.bmm(
@@ -720,19 +555,9 @@ class DeepseekMLAForwardMixin:
q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out)
fuse_rope_for_trtllm_mla = self._fuse_rope_for_trtllm_mla(forward_batch)
skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused()
skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla()
if (
self.rotary_emb is not None
and (not fuse_rope_for_trtllm_mla)
and (not skip_rope_for_dsa_tilelang_fused)
and (not skip_rope_for_aiter_fused_mla)
and (
not _use_aiter
or not _is_gfx95_supported
or self.use_dsa
or self.current_attention_backend == "triton"
)
and not fuse_rope_for_trtllm_mla
# Already applied at the q-prep/indexer overlap fork.
and not self._q8kv8_qprep_overlap_pending
):
@@ -770,7 +595,7 @@ class DeepseekMLAForwardMixin:
dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
mla_prefill_cp = mla_use_prefill_cp(forward_batch)
defer_kv_gather_until_after_rope = _should_defer_dsa_cp_kv_gather(
defer_kv_gather_until_after_rope = should_defer_dsa_cp_kv_gather(
dsa_prefill_cp=dsa_prefill_cp,
fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
)
@@ -796,10 +621,7 @@ class DeepseekMLAForwardMixin:
# all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
if get_parallel().dcp_enabled:
if (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
):
if is_dcp_mla_decode_phase(forward_batch):
if not q_replicate_active:
q_nope_out, q_pe = all_gather_q_for_mla_decode(
q_nope_out=q_nope_out,
@@ -852,168 +674,70 @@ class DeepseekMLAForwardMixin:
save_kv_cache = True
if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
if self._skip_rope_for_dsa_tilelang_fused() and self.rotary_emb is not None:
cos = self.rotary_emb.cos_cache
sin = self.rotary_emb.sin_cache
kv_cache_dtype = (
fp8_dtype if self.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
extra_args = {}
if self._fuse_rope_for_trtllm_mla(forward_batch):
extra_args = {
"cos_sin_cache": self.rotary_emb.cos_sin_cache,
"is_neox": self.rotary_emb.is_neox_style,
"llama_4_scaling": llama_4_scaling,
}
if fusion_plan is not None:
bmm_attention_fn = (
bcg_mla_bmm_then_unified_attention
if is_in_breakable_cuda_graph()
else mla_bmm_then_unified_attention
)
q_cat, _, k_pe_fused, _ = fused_qk_rope_cat_and_cache_mla(
bmm_attention_fn(
fusion_plan.q_nope_t,
self.w_kc,
fusion_plan.q_nope_out_buf,
q_nope_out,
q_pe,
k_nope,
fusion_plan.attn_output_buf,
save_kv_cache,
self.layer_id,
q_pe,
k_pe,
get_token_to_kv_pool().get_key_buffer(self.attn_mqa.layer_id),
forward_batch.out_cache_loc,
positions,
cos,
sin,
self.attn_mqa.k_scale,
self.rotary_emb.is_neox_style,
q_out_dtype=kv_cache_dtype,
cos_sin_cache=extra_args.get("cos_sin_cache"),
is_neox=extra_args.get("is_neox"),
llama_4_scaling=extra_args.get("llama_4_scaling"),
topk_indices=topk_indices,
)
attn_output = fusion_plan.attn_output_buf
elif is_dcp_mla_decode_phase(forward_batch):
# set return_lse=True to correct attn_output
attn_output, lse = self.attn_mqa_for_dcp_decode(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**extra_args,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
save_kv_cache = False
# On decode, pass q_cat directly to attn_mqa with q_rope=None so
# dsa_backend.forward_decode reuses q_cat as a zero-copy view
# (`q.contiguous().view(...)` fast-path) instead of running the
# redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)`
# that would otherwise rebuild a tensor byte-identical to q_cat.
# On ROCm tilelang decode, this eliminates the
# `CatArrayBatchedCopy<OpaqueType<1u>, ...>` kernel that used to
# fire once per layer per decode step (~2.6 us / layer saved).
# Prefill keeps the split form because dsa_backend.forward_extend
# asserts `q_rope is not None`.
if forward_batch.forward_mode.is_decode_or_idle():
if llama_4_scaling is not None:
# llama_4_scaling applies only to the q_nope portion;
# mutate in place via the slice view of q_cat.
q_cat[..., : self.kv_lora_rank] *= llama_4_scaling
attn_output = self.attn_mqa(
q_cat,
None,
None,
forward_batch,
q_rope=None,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
q_nope_fused = q_cat[..., : self.kv_lora_rank]
q_pe_fused = q_cat[..., self.kv_lora_rank :]
if llama_4_scaling is not None:
q_nope_fused *= llama_4_scaling
attn_output = self.attn_mqa(
q_nope_fused,
None,
None,
forward_batch,
q_rope=q_pe_fused,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
extra_args = {}
if self._fuse_rope_for_trtllm_mla(forward_batch):
extra_args = {
"cos_sin_cache": self.rotary_emb.cos_sin_cache,
"is_neox": self.rotary_emb.is_neox_style,
"llama_4_scaling": llama_4_scaling,
}
if fusion_plan is not None:
bmm_attention_fn = (
bcg_mla_bmm_then_unified_attention
if is_in_breakable_cuda_graph()
else mla_bmm_then_unified_attention
)
bmm_attention_fn(
fusion_plan.q_nope_t,
self.w_kc,
fusion_plan.q_nope_out_buf,
q_nope_out,
k_nope,
fusion_plan.attn_output_buf,
save_kv_cache,
self.layer_id,
q_pe,
k_pe,
cos_sin_cache=extra_args.get("cos_sin_cache"),
is_neox=extra_args.get("is_neox"),
llama_4_scaling=extra_args.get("llama_4_scaling"),
topk_indices=topk_indices,
)
attn_output = fusion_plan.attn_output_buf
elif (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
) and get_parallel().dcp_enabled:
# set return_lse=True to correct attn_output
attn_output, lse = self.attn_mqa_for_dcp_decode(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**extra_args,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
attn_output = self.attn_mqa(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**extra_args,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
attn_output = self.attn_mqa(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**extra_args,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
if _use_aiter_gfx95 and self.current_attention_backend == "aiter":
cos = self.rotary_emb.cos_cache
sin = self.rotary_emb.sin_cache
kv_cache_dtype = (
fp8_dtype if self.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
)
q, _, _, k = fused_qk_rope_cat_and_cache_mla(
q_nope_out,
q_pe,
k_nope,
k_pe,
get_token_to_kv_pool().get_key_buffer(self.attn_mqa.layer_id),
forward_batch.out_cache_loc,
positions,
cos,
sin,
self.attn_mqa.k_scale,
self.rotary_emb.is_neox_style,
q_out_dtype=kv_cache_dtype,
)
save_kv_cache = False
else:
q = torch.cat([q_nope_out, q_pe], dim=-1)
k = torch.cat([k_nope, k_pe], dim=-1)
q = torch.cat([q_nope_out, q_pe], dim=-1)
k = torch.cat([k_nope, k_pe], dim=-1)
# Apply llama 4 scaling if provided
if llama_4_scaling is not None:
@@ -1029,10 +753,7 @@ class DeepseekMLAForwardMixin:
)
# correct attn_output with respect to lse from other ranks
if (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
) and get_parallel().dcp_enabled:
if is_dcp_mla_decode_phase(forward_batch):
attn_output = attn_output.view(
-1,
self.num_local_heads * get_parallel().attn_dcp_size,
@@ -1047,7 +768,7 @@ class DeepseekMLAForwardMixin:
)
else:
dcp_comm_backend = get_parallel().dcp_comm_backend
is_lse_base_on_e = _is_mla_dcp_lse_base_on_e(
is_lse_base_on_e = is_mla_dcp_lse_base_on_e(
self.current_attention_backend
)
if dcp_comm_backend in ("a2a", "fi_a2a"):
@@ -1101,83 +822,6 @@ class DeepseekMLAForwardMixin:
attn_bmm_output = (
attn_bmm_output[:, :expected_m, :].transpose(0, 1).flatten(1, 2)
)
elif _is_hip:
# TODO(haishaw): add bmm_fp8 to ROCm
if _use_aiter_gfx95 and self.w_vc.dtype == torch.uint8:
x = attn_output.transpose(0, 1)
B_heads, M_batch = x.shape[0], x.shape[1]
N_vdim = self.w_vc.shape[2]
# Allocate in (batch, heads, dim) so the post-GEMM
# transpose+flatten is a free view instead of a copy.
_bmm_buf = torch.empty(
M_batch,
B_heads,
N_vdim,
device=x.device,
dtype=torch.bfloat16,
)
attn_bmm_output = _bmm_buf.transpose(0, 1)
batched_gemm_afp4wfp4_pre_quant(
x,
self.w_vc.transpose(-2, -1),
self.w_scale_v.transpose(-2, -1),
torch.bfloat16,
attn_bmm_output,
)
else:
_bmm_buf = None
if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn:
attn_bmm_output = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
X=attn_output,
WQ=self.w_vc.transpose(-1, -2),
w_scale=self.w_scale,
group_size=128,
YQ=None,
transpose_bm=False,
transpose_bm_in=True,
dtype=torch.bfloat16,
)
else:
attn_bmm_output = torch.bmm(
attn_output.to(torch.bfloat16).transpose(0, 1),
self.w_vc.to(torch.bfloat16) * self.w_scale,
)
if _bmm_buf is not None:
# _bmm_buf is already (batch, heads, dim) contiguous
if self.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
elif self.o_proj.weight.dtype == torch.float8_e4m3fn:
attn_bmm_output = fused_flatten_fp8_group_quant(
_bmm_buf,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
attn_bmm_output
)
else:
attn_bmm_output = _bmm_buf.flatten(1, 2)
elif self.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
elif self.o_proj.weight.dtype == torch.float8_e4m3fn:
attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_fp8_group_quant(
attn_bmm_output,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
attn_bmm_output
)
else:
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
elif self.w_vc.dtype == torch.float8_e4m3fn:
if _is_cpu:
attn_bmm_output = torch.bmm(
@@ -1276,27 +920,6 @@ class DeepseekMLAForwardMixin:
and get_attn_backend().data_type == torch.float8_e4m3fn
)
def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
"""
Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95.
"""
return (
_use_aiter_gfx95
and self.current_attention_backend in ("dsa", "nsa")
and (
get_exec().kernel.dsa_decode_backend == "tilelang"
or get_exec().kernel.dsa_prefill_backend == "tilelang"
)
)
def _skip_rope_for_aiter_fused_mla(self: DeepseekV2AttentionMLA) -> bool:
"""
Skip rope in prepare and let the fused kernel in forward_absorb_core handle it,
when running aiter-backend MLA on gfx95 (i.e., the `else` branch in forward_absorb_core
that calls fused_qk_rope_cat_and_cache_mla).
"""
return _use_aiter_gfx95 and self.current_attention_backend == "aiter"
# Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output`
# into one eager split op under both PCG and BCG. Without this, the bf16
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py
+1 -1
View File
@@ -29,7 +29,7 @@ if _is_hip:
)
class DeepseekMLARocmForwardMixin:
class DeepseekMLAFusedRopeRocmForwardMixin:
def init_mla_fused_rope_rocm_forward(self: DeepseekV2AttentionMLA):
self.rocm_fused_decode_mla = get_bool_env_var(
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py
+849
View File
@@ -0,0 +1,849 @@
"""AMD/ROCm absorbed multi-latent attention forward path for DeepSeek models.
`AttnForwardMethod.MLA_ROCM` routes here, which keeps every AITER/gfx95 kernel
choice out of the shared `forward_mla.py` and lets non-AMD builds avoid
importing `aiter` altogether.
The BMM absorb steps stay module level to keep ROCm kernel selection localized.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional, Tuple
import torch
from sglang.kernels.ops.quantization.fp8_kernel import (
fp8_dtype,
per_token_group_quant_mla_deep_gemm_masked_fp8,
)
from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mla_extend,
all_gather_q_for_mla_decode,
cp_lse_ag_out_rs_mla,
dcp_a2a_lse_reduce,
)
from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
from sglang.srt.layers.quantization.fp8_utils import (
materialize_bpreshuffle_fp8_scale_tuple,
)
from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
from sglang.srt.lora.deepseek_mla_correction import (
apply_q_correction as apply_kv_b_lora_q_correction,
)
from sglang.srt.lora.deepseek_mla_correction import (
apply_v_correction as apply_kv_b_lora_v_correction,
)
from sglang.srt.lora.deepseek_mla_correction import (
is_kv_b_lora_active,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
_select_local_dcp_heads_for_autotune,
is_dcp_mla_decode_phase,
is_mla_dcp_lse_base_on_e,
should_defer_dsa_cp_kv_gather,
)
from sglang.srt.models.deepseek_common.utils import (
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
_is_gfx95_supported,
_use_aiter,
_use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95,
)
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.state_capturer.indexer_topk import (
maybe_capture_indexer_topk,
)
from sglang.srt.utils import BumpAllocator
logger = logging.getLogger(__name__)
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
if TYPE_CHECKING:
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
if _use_aiter:
# aiter ROCm/aiter#2958 renamed the public `fused_qk_rmsnorm` in
# `aiter.ops.fused_qk_norm_rope_cache_quant` to a private `_fused_qk_rmsnorm`
# and introduced a unified entry point in `aiter.ops.fused_qk_rmsnorm_group_quant`
# with a different (in-place, kwarg-only, no-return) signature. Probe for the
# new symbol first so SGLang works with both pre- and post-#2958 aiter without
# requiring the docker pin to be bumped atomically.
try:
from aiter.ops.enum import QuantType as _AiterQuantType
from aiter.ops.fused_qk_rmsnorm_group_quant import (
fused_qk_rmsnorm as _aiter_fused_qk_rmsnorm_unified,
)
def fused_qk_rmsnorm_bf16(q, q_weight, q_eps, k, k_weight, k_eps):
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
_aiter_fused_qk_rmsnorm_unified(
q_out_quantized=q_out,
k_out=k_out,
q=q,
q_weight=q_weight,
q_epsilon=q_eps,
k=k,
k_weight=k_weight,
k_epsilon=k_eps,
quant_type=_AiterQuantType.No,
)
return q_out, k_out
except ImportError:
from aiter.ops.fused_qk_norm_rope_cache_quant import (
fused_qk_rmsnorm as fused_qk_rmsnorm_bf16,
)
from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
)
if _use_aiter_gfx95:
from aiter.ops.triton.fused_fp8_quant import (
fused_flatten_fp8_group_quant,
fused_rms_fp8_group_quant,
)
from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
batched_gemm_afp4wfp4_pre_quant,
fused_flatten_mxfp4_quant,
fused_rms_mxfp4_quant,
)
from sglang.srt.layers.rocm_linear_utils import fused_qk_rope_cat_and_cache_mla
def rocm_absorb_q_bmm(
attn: DeepseekV2AttentionMLA,
q_nope: torch.Tensor,
*,
is_capture_mode: bool,
) -> torch.Tensor:
"""Absorb ``q_nope @ w_kc`` on HIP/AITER (pre-transpose layout)."""
# TODO(haishaw): add bmm_fp8 to ROCm
if _use_aiter_gfx95 and attn.w_kc.dtype == torch.uint8:
x = q_nope.transpose(0, 1)
q_nope_out = torch.empty(
x.shape[0],
x.shape[1],
attn.w_kc.shape[2],
device=x.device,
dtype=torch.bfloat16,
)
batched_gemm_afp4wfp4_pre_quant(
x,
attn.w_kc.transpose(-2, -1),
attn.w_scale_k.transpose(-2, -1),
torch.bfloat16,
q_nope_out,
)
else:
if (_use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn) or (
is_capture_mode and attn.w_kc.dtype == torch.float8_e4m3fnuz
):
# fp8 Triton kernel: always on gfx950,
# cudagraph-only on gfx942 (hides launch overhead)
q_nope_out = (
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
X=q_nope,
WQ=attn.w_kc.transpose(-1, -2),
w_scale=attn.w_scale,
group_size=128,
YQ=None, # allocate (B, M, N)
transpose_bm=False, # (B, M, N)
transpose_bm_in=True, # (M, B, K)
dtype=torch.bfloat16,
)
)
else:
q_nope_out = torch.bmm(
q_nope.to(torch.bfloat16).transpose(0, 1),
attn.w_kc.to(torch.bfloat16) * attn.w_scale,
)
return q_nope_out
def rocm_absorb_v_bmm(
attn: DeepseekV2AttentionMLA,
attn_output: torch.Tensor,
) -> torch.Tensor:
"""Absorb ``attn_output @ w_vc`` (+ optional fused flatten quant) on HIP."""
# TODO(haishaw): add bmm_fp8 to ROCm
if _use_aiter_gfx95 and attn.w_vc.dtype == torch.uint8:
x = attn_output.transpose(0, 1)
B_heads, M_batch = x.shape[0], x.shape[1]
N_vdim = attn.w_vc.shape[2]
# Allocate in (batch, heads, dim) so the post-GEMM
# transpose+flatten is a free view instead of a copy.
_bmm_buf = torch.empty(
M_batch,
B_heads,
N_vdim,
device=x.device,
dtype=torch.bfloat16,
)
attn_bmm_output = _bmm_buf.transpose(0, 1)
batched_gemm_afp4wfp4_pre_quant(
x,
attn.w_vc.transpose(-2, -1),
attn.w_scale_v.transpose(-2, -1),
torch.bfloat16,
attn_bmm_output,
)
else:
_bmm_buf = None
if _use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn:
attn_bmm_output = (
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
X=attn_output,
WQ=attn.w_vc.transpose(-1, -2),
w_scale=attn.w_scale,
group_size=128,
YQ=None,
transpose_bm=False,
transpose_bm_in=True,
dtype=torch.bfloat16,
)
)
else:
attn_bmm_output = torch.bmm(
attn_output.to(torch.bfloat16).transpose(0, 1),
attn.w_vc.to(torch.bfloat16) * attn.w_scale,
)
if _bmm_buf is not None:
# _bmm_buf is already (batch, heads, dim) contiguous
if attn.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
attn_bmm_output = fused_flatten_fp8_group_quant(
_bmm_buf,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
attn_bmm_output
)
else:
attn_bmm_output = _bmm_buf.flatten(1, 2)
elif attn.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_fp8_group_quant(
attn_bmm_output,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(attn_bmm_output)
else:
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
return attn_bmm_output
def _fused_rope_cat_and_cache(
attn: DeepseekV2AttentionMLA,
q_nope_out: torch.Tensor,
q_pe: torch.Tensor,
k_nope: torch.Tensor,
k_pe: torch.Tensor,
positions: torch.Tensor,
out_cache_loc: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""RoPE + concat + KV-cache write via the AITER fused kernel on gfx95."""
kv_cache_dtype = (
fp8_dtype if attn.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
)
return fused_qk_rope_cat_and_cache_mla(
q_nope_out,
q_pe,
k_nope,
k_pe,
get_token_to_kv_pool().get_key_buffer(attn.attn_mqa.layer_id),
out_cache_loc,
positions,
attn.rotary_emb.cos_cache,
attn.rotary_emb.sin_cache,
attn.attn_mqa.k_scale,
attn.rotary_emb.is_neox_style,
q_out_dtype=kv_cache_dtype,
)
class DeepseekMLARocmForwardMixin:
def forward_absorb_rocm_prepare(
self: DeepseekV2AttentionMLA,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
llama_4_scaling: Optional[torch.Tensor] = None,
prev_topk_indices: Optional[torch.Tensor] = None,
):
from sglang.srt.model_executor.runner import get_is_capture_mode
q_replicate_active = (
get_parallel().dcp_replicate_q_proj
and is_dcp_mla_decode_phase(forward_batch)
and not self.use_deep_gemm_bmm
and self.w_kc_qrep is not None
and self.q_b_proj_qrep_weight is not None
)
q_lora = None
topk_indices = None
if self.q_lora_rank is not None:
q, latent_cache = (
get_attn_tp_context()
.fetch_qkv_latent()
.split(
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
dim=-1,
)
)
k_nope = latent_cache[..., : self.kv_lora_rank]
# overlap qk norm
if self.alt_stream is not None and get_is_capture_mode():
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
q = self.q_a_layernorm(q)
with torch.cuda.stream(self.alt_stream):
k_nope = self.kv_a_layernorm(k_nope)
current_stream.wait_stream(self.alt_stream)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
q, _, k_nope, *_ = fused_rms_mxfp4_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
if self.use_dsa:
q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q_quanted = materialize_bpreshuffle_fp8_scale_tuple(q_quanted)
q = q_quanted
else:
q, _, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
transpose_scale=False,
)
if _use_aiter_bpreshuffle_gfx95:
q = materialize_bpreshuffle_fp8_scale_tuple(q)
elif _use_aiter:
q, k_nope = fused_qk_rmsnorm_bf16(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
# q_lora needed by indexer
if self.use_dsa:
if q_lora is None:
q_lora = q
# overlap q_b_proj and indexer during decode
if (
self.alt_stream is not None
and get_is_capture_mode()
and forward_batch.forward_mode.is_decode_or_idle()
and q_lora is not None
and not q_replicate_active
):
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
with torch.cuda.stream(self.alt_stream):
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj_forward(q)
if self.should_run_indexer(prev_topk_indices):
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
else:
# skip_topk reuses prev layer's indices; mirror into this
# layer's slot so the captured buffer matches what's used.
topk_indices = maybe_capture_indexer_topk(
self.layer_id, prev_topk_indices
)
current_stream.wait_stream(self.alt_stream)
else:
k_nope = k_nope.unsqueeze(1)
if q_replicate_active:
q = torch.nn.functional.linear(q, self.q_b_proj_qrep_weight).view(
-1,
self.num_local_heads * get_parallel().attn_dcp_size,
self.qk_head_dim,
)
else:
q = self.q_b_proj_forward(q)
if q_lora is not None:
if self.should_run_indexer(prev_topk_indices):
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
else:
topk_indices = maybe_capture_indexer_topk(
self.layer_id, prev_topk_indices
)
else:
if q_replicate_active:
q = torch.nn.functional.linear(
hidden_states, self.q_b_proj_qrep_weight
).view(
-1,
self.num_local_heads * get_parallel().attn_dcp_size,
self.qk_head_dim,
)
else:
q = self.q_proj(hidden_states)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
k_nope = latent_cache[..., : self.kv_lora_rank]
k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1)
q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache)
if q_replicate_active:
q_nope_out = (
torch.bmm(q_nope.transpose(0, 1), self.w_kc_qrep)
.transpose(0, 1)
.contiguous()
)
else:
_kvb_q = None
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
# Fork the kv_b q-correction A-step onto the LoRA side stream to overlap the bmm.
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
kv_b_lora_q_prepare,
)
_kvb_q = kv_b_lora_q_prepare(self, q_nope)
if self.use_deep_gemm_bmm:
(
q_nope_val,
q_nope_scale,
masked_m,
expected_m,
aligned_m,
) = per_token_group_quant_mla_deep_gemm_masked_fp8(
q_nope.transpose(0, 1)
)
q_nope_out = q_nope.new_empty(
(self.num_local_heads, aligned_m, self.kv_lora_rank)
)
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
(q_nope_val, q_nope_scale),
(self.w_kc, self.w_scale_k),
q_nope_out,
masked_m,
expected_m,
)
q_nope_out = q_nope_out[:, :expected_m, :]
else:
q_nope_out = rocm_absorb_q_bmm(
self, q_nope, is_capture_mode=get_is_capture_mode()
)
q_nope_out = q_nope_out.transpose(0, 1)
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
kv_b_lora_q_apply,
)
q_nope_out = kv_b_lora_q_apply(self, q_nope, q_nope_out, _kvb_q)
elif is_kv_b_lora_active(self):
q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out)
fuse_rope_for_trtllm_mla = self._fuse_rope_for_trtllm_mla(forward_batch)
if (
self.rotary_emb is not None
and (not fuse_rope_for_trtllm_mla)
and (not self._skip_rope_for_dsa_tilelang_fused())
and (not self._skip_rope_for_aiter_fused_mla())
and (
not _use_aiter
or not _is_gfx95_supported
or self.use_dsa
or self.current_attention_backend == "triton"
)
):
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
mla_prefill_cp = mla_use_prefill_cp(forward_batch)
defer_kv_gather_until_after_rope = should_defer_dsa_cp_kv_gather(
dsa_prefill_cp=dsa_prefill_cp,
fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
)
if dsa_prefill_cp and not defer_kv_gather_until_after_rope:
from sglang.srt.layers.attention.dsa_backend import materialize_full_kv_cp
k_nope, k_pe = materialize_full_kv_cp(
self,
forward_batch,
latent_cache,
k_nope,
k_pe,
)
elif mla_prefill_cp and not is_cp_v2_active(forward_batch):
# CP-v1 gathers the latent here; CP-v2 gathers it in the attention
# backend via the strategy (materialize_full_mla_kv).
k_nope, k_pe = self.rebuild_cp_kv_cache(
latent_cache,
forward_batch,
k_nope,
k_pe,
)
# all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
if get_parallel().dcp_enabled:
if is_dcp_mla_decode_phase(forward_batch):
if not q_replicate_active:
q_nope_out, q_pe = all_gather_q_for_mla_decode(
q_nope_out=q_nope_out,
q_pe=q_pe,
)
elif forward_batch.forward_mode.is_extend():
# for extend, gather kv
all_gather_kv_cache_for_mla_extend(
get_token_to_kv_pool(),
self.attn_mqa,
forward_batch.extend_prefix_lens_cpu,
forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices,
forward_batch.attn_dcp_metadata.dcp_extend_prefix_lens_sum,
forward_batch.attn_dcp_metadata.dcp_kv_buffer,
self.kv_lora_rank,
k_nope,
k_pe,
)
else:
logger.warning(
f"not supported forward_mode {forward_batch.forward_mode}"
)
return (
q_pe,
k_pe,
q_nope_out,
k_nope,
forward_batch,
zero_allocator,
positions,
topk_indices,
llama_4_scaling,
)
def forward_absorb_rocm_core(
self: DeepseekV2AttentionMLA,
q_pe,
k_pe,
q_nope_out,
k_nope,
forward_batch,
zero_allocator,
positions,
topk_indices,
llama_4_scaling,
):
save_kv_cache = True
if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
if self._skip_rope_for_dsa_tilelang_fused() and self.rotary_emb is not None:
q_cat, _, k_pe_fused, _ = _fused_rope_cat_and_cache(
self,
q_nope_out,
q_pe,
k_nope,
k_pe,
positions,
forward_batch.out_cache_loc,
)
save_kv_cache = False
# On decode, pass q_cat directly to attn_mqa with q_rope=None so
# dsa_backend.forward_decode reuses q_cat as a zero-copy view
# (`q.contiguous().view(...)` fast-path) instead of running the
# redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)`
# that would otherwise rebuild a tensor byte-identical to q_cat.
# On ROCm tilelang decode, this eliminates the
# `CatArrayBatchedCopy<OpaqueType<1u>, ...>` kernel that used to
# fire once per layer per decode step (~2.6 us / layer saved).
# Prefill keeps the split form because dsa_backend.forward_extend
# asserts `q_rope is not None`.
if forward_batch.forward_mode.is_decode_or_idle():
if llama_4_scaling is not None:
# llama_4_scaling applies only to the q_nope portion;
# mutate in place via the slice view of q_cat.
q_cat[..., : self.kv_lora_rank] *= llama_4_scaling
attn_output = self.attn_mqa(
q_cat,
None,
None,
forward_batch,
q_rope=None,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
q_nope_fused = q_cat[..., : self.kv_lora_rank]
q_pe_fused = q_cat[..., self.kv_lora_rank :]
if llama_4_scaling is not None:
q_nope_fused *= llama_4_scaling
attn_output = self.attn_mqa(
q_nope_fused,
None,
None,
forward_batch,
q_rope=q_pe_fused,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
extra_args = {}
if self._fuse_rope_for_trtllm_mla(forward_batch):
extra_args = {
"cos_sin_cache": self.rotary_emb.cos_sin_cache,
"is_neox": self.rotary_emb.is_neox_style,
"llama_4_scaling": llama_4_scaling,
}
if is_dcp_mla_decode_phase(forward_batch):
# set return_lse=True to correct attn_output
attn_output, lse = self.attn_mqa_for_dcp_decode(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**extra_args,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
attn_output = self.attn_mqa(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**extra_args,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
if self._skip_rope_for_aiter_fused_mla():
q, _, _, k = _fused_rope_cat_and_cache(
self,
q_nope_out,
q_pe,
k_nope,
k_pe,
positions,
forward_batch.out_cache_loc,
)
save_kv_cache = False
else:
q = torch.cat([q_nope_out, q_pe], dim=-1)
k = torch.cat([k_nope, k_pe], dim=-1)
# Apply llama 4 scaling if provided
if llama_4_scaling is not None:
q *= llama_4_scaling
attn_output = self.attn_mqa(
q,
k,
k_nope,
forward_batch,
save_kv_cache=save_kv_cache,
**(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
)
# correct attn_output with respect to lse from other ranks
if is_dcp_mla_decode_phase(forward_batch):
attn_output = attn_output.view(
-1,
self.num_local_heads * get_parallel().attn_dcp_size,
self.kv_lora_rank,
)
if get_in_autotune_dummy_run():
# The synthetic FlashInfer MoE autotune pass discards model
# outputs. Avoid an unnecessary cross-node MNNVL exchange of
# zero attention partials.
attn_output = _select_local_dcp_heads_for_autotune(
attn_output, self.num_local_heads
)
else:
dcp_comm_backend = get_parallel().dcp_comm_backend
is_lse_base_on_e = is_mla_dcp_lse_base_on_e(
self.current_attention_backend
)
if dcp_comm_backend in ("a2a", "fi_a2a"):
# A2A exchange of head partials + LSE, then local Triton combine.
attn_output = dcp_a2a_lse_reduce(
attn_output.contiguous(),
lse.contiguous(),
get_parallel().dcp_group,
is_lse_base_on_e=is_lse_base_on_e,
comm_backend=dcp_comm_backend,
)
else:
attn_output = cp_lse_ag_out_rs_mla(
attn_output,
lse,
get_parallel().dcp_group,
is_lse_base_on_e=is_lse_base_on_e,
)
attn_output = attn_output.transpose(0, 1)
attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
_kvb_v = None
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
# Fork the kv_b v-correction A-step onto the LoRA side stream to overlap the bmm.
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
kv_b_lora_v_prepare,
)
_kvb_v = kv_b_lora_v_prepare(self, attn_output)
if self.use_deep_gemm_bmm:
(
attn_output_val,
attn_output_scale,
masked_m,
expected_m,
aligned_m,
) = per_token_group_quant_mla_deep_gemm_masked_fp8(
attn_output.transpose(0, 1)
)
attn_bmm_output = attn_output.new_empty(
(self.num_local_heads, aligned_m, self.v_head_dim)
)
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
(attn_output_val, attn_output_scale),
(self.w_vc, self.w_scale_v),
attn_bmm_output,
masked_m,
expected_m,
)
attn_bmm_output = (
attn_bmm_output[:, :expected_m, :].transpose(0, 1).flatten(1, 2)
)
else:
attn_bmm_output = rocm_absorb_v_bmm(self, attn_output)
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
kv_b_lora_v_apply,
)
attn_bmm_output = kv_b_lora_v_apply(
self, attn_output, attn_bmm_output, _kvb_v
)
elif is_kv_b_lora_active(self):
attn_bmm_output = apply_kv_b_lora_v_correction(
self, attn_output, attn_bmm_output
)
output, _ = self.o_proj(attn_bmm_output)
if self.next_skip_topk is None:
return output
# Return topk_indices for the next layer when enabling index cache
if not self.next_skip_topk:
return output, None
else:
return output, topk_indices
def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
"""
Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95.
"""
return (
_use_aiter_gfx95
and self.current_attention_backend in ("dsa", "nsa")
and (
get_exec().kernel.dsa_decode_backend == "tilelang"
or get_exec().kernel.dsa_prefill_backend == "tilelang"
)
)
def _skip_rope_for_aiter_fused_mla(self: DeepseekV2AttentionMLA) -> bool:
"""
Skip rope in prepare and let the fused kernel in forward_absorb_rocm_core handle it,
when running aiter-backend MLA on gfx95 (i.e., the `else` branch in
forward_absorb_rocm_core that calls fused_qk_rope_cat_and_cache_mla).
"""
return _use_aiter_gfx95 and self.current_attention_backend == "aiter"
python/sglang/srt/models/deepseek_v2.py
+29 -1
View File
@@ -162,12 +162,15 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
)
from sglang.srt.models.deepseek_common.attention_backend_handler import (
AttentionBackendRegistry,
resolve_rocm_forward_method,
)
from sglang.srt.models.deepseek_common.attention_forward_methods import (
AttnForwardMethod,
DeepseekMHAForwardMixin,
DeepseekMHARocmForwardMixin,
DeepseekMLACpuForwardMixin,
DeepseekMLAForwardMixin,
DeepseekMLAFusedRopeRocmForwardMixin,
DeepseekMLARocmForwardMixin,
)
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
@@ -1711,8 +1714,10 @@ class DeepseekV2MoE(nn.Module):
class DeepseekV2AttentionMLA(
nn.Module,
DeepseekMHAForwardMixin,
DeepseekMHARocmForwardMixin,
DeepseekMLAForwardMixin,
DeepseekMLARocmForwardMixin,
DeepseekMLAFusedRopeRocmForwardMixin,
DeepseekMLACpuForwardMixin,
):
@@ -2021,7 +2026,7 @@ class DeepseekV2AttentionMLA(
self.current_attention_backend = attention_backend
handler = AttentionBackendRegistry.get_handler(attention_backend)
return handler(self, forward_batch)
return resolve_rocm_forward_method(handler(self, forward_batch))
def op_prepare(self, state):
state.attn_intermediate_state = self.forward_prepare(
@@ -2117,6 +2122,23 @@ class DeepseekV2AttentionMLA(
llama_4_scaling,
prev_topk_indices,
)
elif attn_forward_method == AttnForwardMethod.MHA_ROCM:
inner_state = self.forward_normal_rocm_prepare(
positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT_ROCM:
inner_state = self.forward_normal_one_shot_rocm_prepare(
positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.MLA_ROCM:
inner_state = self.forward_absorb_rocm_prepare(
positions,
hidden_states,
forward_batch,
zero_allocator,
llama_4_scaling,
prev_topk_indices,
)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
inner_state = self.forward_absorb_fused_mla_rope_prepare(
positions, hidden_states, forward_batch, zero_allocator
@@ -2172,6 +2194,12 @@ class DeepseekV2AttentionMLA(
return self.forward_normal_one_shot_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA:
return self.forward_absorb_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MHA_ROCM:
return self.forward_normal_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT_ROCM:
return self.forward_normal_one_shot_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA_ROCM:
return self.forward_absorb_rocm_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
return self.forward_absorb_fused_mla_rope_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
test/registered/kernels/test_dcp_lse_combine.py
+3 -3
View File
@@ -267,11 +267,11 @@ class TestCPUReference(CustomTestCase):
def test_flashmla_selects_natural_log_lse(self):
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
_is_mla_dcp_lse_base_on_e,
is_mla_dcp_lse_base_on_e,
)
self.assertTrue(_is_mla_dcp_lse_base_on_e("flashmla"))
self.assertFalse(_is_mla_dcp_lse_base_on_e("flashinfer_mla"))
self.assertTrue(is_mla_dcp_lse_base_on_e("flashmla"))
self.assertFalse(is_mla_dcp_lse_base_on_e("flashinfer_mla"))
def test_nan_lse_handled(self):
from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu
test/registered/unit/models/test_deepseek_mla_dispatch.py
+33
View File
@@ -62,5 +62,38 @@ class TestDispatchMLASubtype(CustomTestCase):
self.assertEqual(method, AttnForwardMethod.MLA)
class TestResolveRocmForwardMethod(CustomTestCase):
"""The generic MHA/MLA methods must never reach the CUDA forward paths on
ROCm: those were stripped of their AMD branches when the AITER kernels moved
into forward_mha_rocm.py / forward_mla_rocm.py."""
def test_hip_routes_shared_methods_to_rocm(self):
with mock.patch.object(abh, "_is_hip", True):
self.assertEqual(
abh.resolve_rocm_forward_method(AttnForwardMethod.MHA),
AttnForwardMethod.MHA_ROCM,
)
self.assertEqual(
abh.resolve_rocm_forward_method(AttnForwardMethod.MHA_ONE_SHOT),
AttnForwardMethod.MHA_ONE_SHOT_ROCM,
)
self.assertEqual(
abh.resolve_rocm_forward_method(AttnForwardMethod.MLA),
AttnForwardMethod.MLA_ROCM,
)
def test_hip_leaves_platform_specific_methods_alone(self):
with mock.patch.object(abh, "_is_hip", True):
self.assertEqual(
abh.resolve_rocm_forward_method(AttnForwardMethod.MLA_FUSED_ROPE_ROCM),
AttnForwardMethod.MLA_FUSED_ROPE_ROCM,
)
def test_non_hip_is_identity(self):
with mock.patch.object(abh, "_is_hip", False):
for method in AttnForwardMethod:
self.assertEqual(abh.resolve_rocm_forward_method(method), method)
if __name__ == "__main__":
unittest.main()