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

[Feature] add feature mla_ag_after_qlora for dsv3.2 (#19428)

Browse Source 
Co-authored-by: JiaruiChang5268 <changjiarui1@huawei.com>
This commit is contained in:
JiaruiChang5268
2026-03-02 20:00:31 +08:00
committed by GitHub
parent 3f36f27eae
commit b3718982a1
5 changed files with 101 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/environ.py
View File

@@ -312,6 +312,8 @@ class Envs:
SGLANG_NPU_FORWARD_NATIVE_GELUTANH = EnvBool(False)
# Forward native implementation for gemma rms norm for model Skywork-Reward-Gemma-2-27B-v0.2
SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM = EnvBool(False)
# Delay all-gather after qlora for better performance for Deepseek v3.2
SGLANG_USE_AG_AFTER_QLORA = EnvBool(False)
# Quantization
SGLANG_INT4_WEIGHT = EnvBool(False)

38
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
View File

@@ -4,21 +4,24 @@ from typing import TYPE_CHECKING
import torch
import torch_npu
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
NPUFusedMLAPreprocess,
is_fia_nz,
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full
from sglang.srt.layers.attention.nsa.utils import (
cp_split_and_rebuild_position,
nsa_use_prefill_cp,
)
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
from sglang.srt.utils import BumpAllocator
_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
# region MHA
@@ -28,6 +31,7 @@ def forward_mha_prepare_npu(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
layer_scatter_modes,
):
if m.q_lora_rank is not None:
q, latent_cache = (
@@ -55,6 +59,13 @@ def forward_mha_prepare_npu(
else:
q = m.q_a_layernorm(q)
if (
_use_ag_after_qlora
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
):
q = scattered_to_tp_attn_full(q, forward_batch)
latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
else:
@@ -142,6 +153,7 @@ def forward_mla_prepare_npu(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
layer_scatter_modes,
):
if is_mla_preprocess_enabled():
if not hasattr(m, "mla_preprocess"):
@@ -184,6 +196,13 @@ def forward_mla_prepare_npu(
k_nope = latent_cache[..., : m.kv_lora_rank]
q = m.q_a_layernorm(q)
if (
_use_ag_after_qlora
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
):
q = scattered_to_tp_attn_full(q, forward_batch)
latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
k_nope = m.kv_a_layernorm(k_nope)
# q_lora needed by indexer
@@ -285,6 +304,7 @@ def forward_dsa_prepare_npu(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
layer_scatter_modes,
):
dynamic_scale = None
if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode():
@@ -313,7 +333,13 @@ def forward_dsa_prepare_npu(
# overlap qk norm
q = m.q_a_layernorm(q)
if (
_use_ag_after_qlora
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
):
q = scattered_to_tp_attn_full(q, forward_batch)
latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
q_lora = q.clone() # required for topk_indices
q_event = None
@@ -353,7 +379,13 @@ def forward_dsa_prepare_npu(
)
topk_indices = m.indexer(
hidden_states, q_lora, positions, forward_batch, m.layer_id, dynamic_scale
hidden_states,
q_lora,
positions,
forward_batch,
m.layer_id,
layer_scatter_modes,
dynamic_scale,
)
return (

35
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -12,6 +12,7 @@ from sglang.jit_kernel.fused_store_index_cache import (
fused_store_index_k_cache,
)
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.utils import MultiPlatformOp
@@ -43,6 +44,7 @@ from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
)
from sglang.srt.layers.communicator import ScatterMode
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
@@ -50,6 +52,7 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.server_args import get_global_server_args
_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
@@ -1203,6 +1206,7 @@ class Indexer(MultiPlatformOp):
positions: torch.Tensor,
forward_batch: ForwardBatch,
layer_id: int,
layer_scatter_modes=None,
dynamic_scale: torch.Tensor = None,
) -> torch.Tensor:
if forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int is None:
@@ -1223,7 +1227,7 @@ class Indexer(MultiPlatformOp):
cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
bs = x.shape[0]
bs = q_lora.shape[0]
if self.alt_stream is not None:
self.alt_stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(self.alt_stream):
@@ -1276,6 +1280,12 @@ class Indexer(MultiPlatformOp):
k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128]
k = self.k_norm(k_proj)
if (
_use_ag_after_qlora
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
):
k = scattered_to_tp_attn_full(k, forward_batch)
k_pe, k_nope = torch.split(
k,
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
@@ -1356,7 +1366,12 @@ class Indexer(MultiPlatformOp):
torch.npu.current_stream().wait_event(q_rope_event)
if envs.SGLANG_NPU_USE_MULTI_STREAM.get():
torch.npu.current_stream().wait_event(weights_event)
if (
_use_ag_after_qlora
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
):
weights = scattered_to_tp_attn_full(weights, forward_batch)
block_table = forward_batch.attn_backend.forward_metadata.block_tables
if (
is_prefill
@@ -1450,3 +1465,19 @@ class Indexer(MultiPlatformOp):
sparse_mode=3,
)
return topk_indices_prev[0], topk_indices_next[0]
def scattered_to_tp_attn_full(
hidden_states: torch.Tensor,
forward_batch,
) -> torch.Tensor:
hidden_states, local_hidden_states = (
torch.empty(
(forward_batch.input_ids.shape[0], hidden_states.shape[1]),
dtype=hidden_states.dtype,
device=hidden_states.device,
),
hidden_states,
)
attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states.contiguous())
return hidden_states

6
python/sglang/srt/layers/communicator.py
View File

@@ -29,6 +29,7 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
@@ -74,6 +75,7 @@ _is_sm100_supported = _is_cuda and is_sm100_supported()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
_is_gfx95_supported = is_gfx95_supported()
_is_npu = is_npu()
_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
if _use_aiter and _is_gfx95_supported:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
@@ -189,7 +191,7 @@ class AttnTpContext:
def init_context(self, q_lora_rank, is_nsa):
self.allow_input_scattered = (
get_global_server_args().enable_attn_tp_input_scattered
and _is_cuda
and (_is_cuda or _is_npu)
and q_lora_rank is not None
and not is_nsa
and get_tensor_model_parallel_world_size() > 1
@@ -693,6 +695,8 @@ class CommunicateSimpleFn:
if (input_mode == ScatterMode.SCATTERED) and (
output_mode == ScatterMode.TP_ATTN_FULL
):
if _use_ag_after_qlora:
return CommunicateSimpleFn._trivial
return CommunicateSimpleFn._scattered_to_tp_attn_full
raise NotImplementedError(f"{input_mode=} {output_mode=}")

29
python/sglang/srt/models/deepseek_v2.py
View File

@@ -1301,6 +1301,7 @@ class DeepseekV2AttentionMLA(
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
layer_scatter_modes: LayerScatterModes = None,
llama_4_scaling: Optional[torch.Tensor] = None,
):
s = self.forward_prepare(
@@ -1308,6 +1309,7 @@ class DeepseekV2AttentionMLA(
hidden_states=hidden_states,
forward_batch=forward_batch,
zero_allocator=zero_allocator,
layer_scatter_modes=layer_scatter_modes,
llama_4_scaling=llama_4_scaling,
)
return self.forward_core(s)
@@ -1318,6 +1320,7 @@ class DeepseekV2AttentionMLA(
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
layer_scatter_modes: LayerScatterModes = None,
llama_4_scaling: Optional[torch.Tensor] = None,
):
if self.attn_mha.kv_b_proj is None:
@@ -1370,15 +1373,30 @@ class DeepseekV2AttentionMLA(
)
elif attn_forward_method == AttnForwardMethod.MHA_NPU:
inner_state = forward_mha_prepare_npu(
self, positions, hidden_states, forward_batch, zero_allocator
self,
positions,
hidden_states,
forward_batch,
zero_allocator,
layer_scatter_modes,
)
elif attn_forward_method == AttnForwardMethod.MLA_NPU:
inner_state = forward_mla_prepare_npu(
self, positions, hidden_states, forward_batch, zero_allocator
self,
positions,
hidden_states,
forward_batch,
zero_allocator,
layer_scatter_modes,
)
elif attn_forward_method == AttnForwardMethod.DSA_NPU:
inner_state = forward_dsa_prepare_npu(
self, positions, hidden_states, forward_batch, zero_allocator
self,
positions,
hidden_states,
forward_batch,
zero_allocator,
layer_scatter_modes,
)
else:
raise NotImplementedError
@@ -1505,6 +1523,10 @@ class DeepseekV2DecoderLayer(nn.Module):
prefix=add_prefix("self_attn", prefix),
alt_stream=alt_stream,
)
if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get():
raise ValueError(
"SGLANG_USE_AG_AFTER_QLORA only supports the model with q_lora_rank"
)
self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn)
is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False)
@@ -1627,6 +1649,7 @@ class DeepseekV2DecoderLayer(nn.Module):
forward_batch=forward_batch,
zero_allocator=zero_allocator,
llama_4_scaling=llama_4_scaling,
layer_scatter_modes=self.layer_scatter_modes,
)
hidden_states, residual = self.layer_communicator.prepare_mlp(