diff --git a/python/sglang/srt/batch_overlap/operations.py b/python/sglang/srt/batch_overlap/operations.py index 2b853e4ec..df3decc7d 100644 --- a/python/sglang/srt/batch_overlap/operations.py +++ b/python/sglang/srt/batch_overlap/operations.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 92fab1bc3..9a2a60120 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -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 = {} diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py index 2a9508fe2..029297660 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py @@ -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", ] diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py index c234a6b38..09e6e2a5a 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py @@ -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() diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 19b4805c1..69d90b600 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py new file mode 100644 index 000000000..00fa2407c --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 698d7f047..b3d666f8c 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -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, ...>` 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 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py index 6112b565b..d8bca2075 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py @@ -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( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py new file mode 100644 index 000000000..04f1a8f1b --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -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, ...>` 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" diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 76727bbdf..6fcbc1cb7 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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: diff --git a/test/registered/kernels/test_dcp_lse_combine.py b/test/registered/kernels/test_dcp_lse_combine.py index 3c6eb5086..9d26a6428 100644 --- a/test/registered/kernels/test_dcp_lse_combine.py +++ b/test/registered/kernels/test_dcp_lse_combine.py @@ -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 diff --git a/test/registered/unit/models/test_deepseek_mla_dispatch.py b/test/registered/unit/models/test_deepseek_mla_dispatch.py index 17f6c6378..57571eaa9 100644 --- a/test/registered/unit/models/test_deepseek_mla_dispatch.py +++ b/test/registered/unit/models/test_deepseek_mla_dispatch.py @@ -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()