Motivation

For the classic dense decode layer structure (self-attention + MLP), in the pure TP case, tensors are parallelized in the attention layer and the MLP layer. Since each device contains the full amount of data before and after each TP split, Layernorm stores 2*2BSH bytes of activation value (assume hidden states shape is [B, S, H]). SP aims to split this 4*BSH redundant data across multiple devices, and do layernorm independently. For more details please refer to the research paper https://arxiv.org/pdf/2205.05198.pdf

We add a parameter "--enable-sp" at server args, which can enable sequence parallel if be set. For example

python3 -m sglang.launch_server --model-path $MODEL_PATH \
        --tp-size 16 --dp-size 1 --enable-sp \
        --trust-remote-code --attention-backend ascend --device npu --host $HOST_IP --port $PORT \
        --quantization w8a8_int8 --mem-fraction-static 0.8 \
        --chunked-prefill-size 16000 --context-length 16000 --max-prefill-tokens 16000 --max-total-tokens 16000 \
        --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto

The TP-SP has two benifits:

1. Reduce long dataset (36K) TTFT for 10% on deepseek v3 and 7% on deepseek v3.2.
2. Reduce peak memory because of less activations on RMSNorm layer.

Modifications

1. For the RowParallel linear, we replace all-reduce comm op with reduce-scatter on TP group if enable sp.
2. For the dense layers of model (qwen2/3 & deepseek v2/3/3.2), we split residual before layernorm at the first layer, hidden states is be splitted at RowParallel linear so do not need to split again.
3. After the first dense layer, all data are scattered state, only do layernorm at prepare_attn and prepare_mlp.
4. Before the MLP and ATTENTION module, we do extra all-gather because of weights has been splitted on tensor dimension, to make sure inputs are complete.
5. Do all-gather at the last dense layer.
6. For sparse layer of deepseek models, when enableed Deepep, the hidden states has been sscattered, we utilize it and move the _scattered_to_tp_attn_full from prepare_attention to after getting the q_lora and latent cache, which also can decrease much more computation.

NOTE:

1. We only adapt qwen2/3 and deepseek v2/3/3.2 on ascend backend, for other backends you can add few code to adapt SP.
2. TP-SP can enable with CP ([Ascend] Deepseek v3 and v3.2 support Context Parallelism #12207) together on Ascend backend. We test it on dsv3.2 model, CP=16, TP=2, which can continue reducing TTFT for 1 percent and decrease peak runtime memory (2G for single device).

Accuracy Tests

The test_gsam8k.py is based on benchmark/gsm8k/bench_sglang.py, you can test it with
python3 benchmark/gsm8k/bench_sglang.py --host http://127.0.0.1 --port 8000 --num-questions 300.

Benchmarking and Profiling

export cann_path=/usr/local/Ascend/ascend-toolkit/latest
source /usr/local/Ascend/driver/bin/setenv.bash
source ${cann_path}/../set_env.sh
source ${cann_path}/../../nnal/atb/set_env.sh
source ${cann_path}/opp/vendors/customize/bin/set_env.bash
export ASCEND_HOME_PATH=${cann_path}

# CPU high preformance
echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor
sysctl -w vm.swappiness=0
sysctl -w kernel.numa_balancing=0
sysctl -w kernel.sched_migration_cost_ns=50000

export SGLANG_SET_CPU_AFFINITY=1

# Memory Fragmentation
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export STREAMS_PER_DEVICE=32

# HCCL
export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=16
export HCCL_BUFFSIZE=1600
#export HCCL_RDMA_PCIE_DIRECT_POST_NOSTRICT=TRUE
export HCCL_OP_EXPANSION_MODE=AIV
export HCCL_ALGO="level0:NA;level1:ring"

# Your NIC
export HCCL_SOCKET_IFNAME=enp48s3u1u1
export GLOO_SOCKET_IFNAME=enp48s3u1u1

export PYTHONPATH=$PWD/python/:$PYTHONPATH

python -m sglang.launch_server --model-path $MODEL_PATH \
        --tp-size 16 --dp-size 1 --enable-sp \
        --trust-remote-code --attention-backend ascend --device npu --host 127.0.0.1 --port 8000 \
        --quantization w8a8_int8 --mem-fraction-static 0.79 \
        --chunked-prefill-size 36000 --context-length 36000 --max-prefill-tokens 36000 --max-total-tokens 36000 \
        --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.