Merge remote-tracking branch 'upstream/main' into 0106

* upstream/main:
  fix little batchsize and int8 quant on ci (sgl-project#302)
  optimize sinks attention (sgl-project#260)
  add swiglu_oai_triton (sgl-project#270)
  update tag to 2026.01.12 (sgl-project#312)
  feat:add performance compare (sgl-project#311)
  support add_gemma_rms_norm (sgl-project#310)
  optimize gdn gating and fused_qkvzba_split_reshape_cat (sgl-project#306)
  fix layout numTokensPerExpertTensor partial Initialization bug (sgl-project#303)
  Supplement A2 doc, software and hardware compatibility info (sgl-project#294)
  Added an environment variable to control whether to enable the Combine Ant Migration feature. (sgl-project#304)

Author: zhuyutong332

.github/workflows/pr-test-npu.yml

@@ -74,6 +74,8 @@ jobs:
           HCCL_BUFFSIZE: 1913
         run: |
           python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py
+          python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=1
+          python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=2
 
       - name: Run test base fused deep moe
         timeout-minutes: 10
@@ -168,6 +170,8 @@ jobs:
           HCCL_BUFFSIZE: 1913
         run: |
           python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py
+          python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=1
+          python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=2
 
       - name: Run test base fused deep moe
         timeout-minutes: 10

config.ini

@@ -1,2 +1,2 @@
 [global]
-version = 2025.12.25
+version = 2026.01.12

csrc/deepep/deep_ep.cpp

@@ -48,6 +48,7 @@ Buffer::Buffer(int64_t rank, int64_t num_ranks, int64_t num_nvl_bytes, int64_t n
     this->shared_expert_rank_num = get_value_from_env("MOE_SHARED_EXPERT_RANK_NUM", 0);
     const char *roundEnv = std::getenv("DEEPEP_NORMAL_LONG_SEQ_ROUND");
     const char *tokensEnv = std::getenv("DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS");
+    this->combine_enable_long_seq = get_value_from_env("DEEPEP_NORMAL_COMBINE_ENABLE_LONG_SEQ", 0);
     bool roundSet = (roundEnv != nullptr);
     bool tokensSet = (tokensEnv != nullptr);
 
@@ -602,6 +603,8 @@ Buffer::intranode_combine(const torch::Tensor &x, const torch::Tensor &topk_idx,
     std::optional<torch::Tensor> recv_topk_weights;
     std::optional<EventHandle> event;
 
+    int32_t round = this->combine_enable_long_seq ? this->round : 1;
+    int32_t per_round_tokens = this->combine_enable_long_seq ? this->per_round_tokens : MAX_TOKENS_PER_ROUND;
     EXEC_NPU_CMD(aclnnCamMoeCombineNormal, recv_x, token_src_info, ep_send_counts, expert_scales, tp_send_counts,
                  hcom_ep_name, num_ranks, rank, hcom_ep_name, tp_world_size, tp_rankId, moe_expert_number, real_max_bs,
                  round, per_round_tokens, combined_x, combine_send_cost_stats_out);

csrc/deepep/deep_ep.hpp

@@ -24,6 +24,7 @@ struct Buffer {
 
     int32_t round;
     int32_t per_round_tokens;
+    bool combine_enable_long_seq = false;  // Whether to enable the Combine Ant Migration feature
 
     bool low_latency_mode = false;
     bool is_padding = false;

csrc/deepep/ops/op_kernel/dispatch_layout.h

@@ -127,7 +127,7 @@ class DispatchLayout
             SyncFunc<AscendC::HardEvent::MTE3_V>();
             Duplicate<T>(numTokensPerRankTensor, 0, numRanks_);
             Duplicate<T>(isTokenInRankTensor, 0, tempTokens_ * numRanks_);
-            Duplicate<T>(numTokensPerExpertTensor, 0, numExperts_);
+            Duplicate<T>(numTokensPerExpertTensor, 0, numTokensPerExpert32AlignIntLen_ / sizeof(T));
             SyncFunc<AscendC::HardEvent::V_S>();
             SyncFunc<AscendC::HardEvent::V_MTE3>();
             const DataCopyExtParams clearGmParams{1U, numTokensPerExpert32AlignIntLen_, 0U, 0U, 0U};

csrc/deepep/ops/op_kernel/dispatch_layout_a2.h

@@ -168,8 +168,8 @@ class DispatchLayoutA2
         LocalTensor<T> countExpertTensor = countExpertBuf_.AllocTensor<T>();
         Duplicate<T>(countExpertTensor, 0, numExperts_);
         Duplicate<T>(numTokensPerRankTensor, 0, numRanks_);
-        Duplicate<T>(numTokensPerExpertTensor, 0, numExperts_);
-        Duplicate<T>(prefixCountPerExpertTensor, 0, numExperts_);
+        Duplicate<T>(numTokensPerExpertTensor, 0, numTokensPerExpert32AlignIntLen_ / sizeof(T));
+        Duplicate<T>(prefixCountPerExpertTensor, 0, numTokensPerExpert32AlignIntLen_ / sizeof(T));
         Duplicate<T>(isTokenInRankTensor, 0, tempTokens_ * numRanks_);
         Duplicate<T>(localTokenServerOffsetTensor, 0, localTokenServerOffset32AlignIntLen_ / sizeof(T));
         Duplicate<T>(sendTokenIdxTensor, 0, sendTokenIdx32AlignIntLen_ / sizeof(T));

csrc/deepep/ops2/op_kernel/dispatch_layout.h

@@ -98,7 +98,7 @@ class DispatchLayout
         LocalTensor<T> seenRankTensor = seenRankBuf_.AllocTensor<T>();
         LocalTensor<T> sendTokenIdxSmallTensor = sendTokenIdxSmallBuf_.AllocTensor<T>();
         Duplicate<T>(numTokensPerRankTensor, 0, numRanks_);
-        Duplicate<T>(numTokensPerExpertTensor, 0, numExperts_);
+        Duplicate<T>(numTokensPerExpertTensor, 0, numTokensPerExpert32AlignIntLen_ / sizeof(T));
         Duplicate<T>(isTokenInRankTensor, 0, tempTokens_ * numRanks_);
         SyncFunc<AscendC::HardEvent::V_S>();
 

csrc/deepep/ops2/op_kernel/dispatch_layout_a2.h

@@ -168,8 +168,8 @@ class DispatchLayoutA2
         LocalTensor<T> countExpertTensor = countExpertBuf_.AllocTensor<T>();
         Duplicate<T>(countExpertTensor, 0, numExperts_);
         Duplicate<T>(numTokensPerRankTensor, 0, numRanks_);
-        Duplicate<T>(numTokensPerExpertTensor, 0, numExperts_);
-        Duplicate<T>(prefixCountPerExpertTensor, 0, numExperts_);
+        Duplicate<T>(numTokensPerExpertTensor, 0, numTokensPerExpert32AlignIntLen_ / sizeof(T));
+        Duplicate<T>(prefixCountPerExpertTensor, 0, numTokensPerExpert32AlignIntLen_ / sizeof(T));
         Duplicate<T>(isTokenInRankTensor, 0, tempTokens_ * numRanks_);
         Duplicate<T>(localTokenServerOffsetTensor, 0, localTokenServerOffset32AlignIntLen_ / sizeof(T));
         Duplicate<T>(sendTokenIdxTensor, 0, sendTokenIdx32AlignIntLen_ / sizeof(T));

python/deep_ep/A2_DEEPEP_CN.md

@@ -1,6 +1,12 @@
 
 A2场景下使用DeepEp说明
 
+# 软硬件配套说明
+硬件型号支持：Atlas A2 系列产品
+平台：aarch64/x86
+配套软件
+- 驱动 Ascend HDK ≥ 25.3.RC1、CANN ≥ 8.3.RC1
+
 # 构建DeepEp包
 执行工程构建脚本 build.sh
 ```bash
@@ -47,6 +53,12 @@ DeepEp 向上层提供以下核心接口：
 export HCCL_BUFFSIZE=1024
 ```
 
+A2场景下叠加deepep，需**禁用**环境变量`HCCL_OP_EXPANSION_MODE`，否则会出现未知算子错误。
+```bash
+# A2下需要去除该环境变量
+# export HCCL_OP_EXPANSION_MODE=AIV
+```
+
 ## A2单机
 
 ### 框架接入建议

python/sgl_kernel_npu/sgl_kernel_npu/activation/swiglu_oai.py

@@ -1,9 +1,88 @@
-# This file contains swiglu for OpenAI models.
-# It will be optimized using Triton in the future.
 import torch
+import triton
+import triton.language as tl
+from sgl_kernel_npu.utils.triton_utils import get_device_properties
 
 
-def swiglu_oai(layer, hidden_states):
+@triton.jit
+def swiglu_oai_kernel(
+    hidden_states,
+    gated_output,
+    gemm1_alpha,
+    gemm1_clamp_limit,
+    output_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    MINIBLOCK_SIZE: tl.constexpr,
+    BS: tl.constexpr,
+):
+    i_block = tl.program_id(0)
+
+    for i_miniblock in range(0, BLOCK_SIZE, MINIBLOCK_SIZE):
+        offset_bs = i_block * BLOCK_SIZE + i_miniblock + tl.arange(0, MINIBLOCK_SIZE)
+        mask_bs = offset_bs < BS
+
+        offset_gate = tl.arange(0, output_dim) * 2
+        offset_up = tl.arange(0, output_dim) * 2 + 1
+
+        gate = tl.load(
+            hidden_states + offset_bs[:, None] * output_dim * 2 + offset_gate[None, :],
+            mask=mask_bs[:, None],
+        )
+        up = tl.load(
+            hidden_states + offset_bs[:, None] * output_dim * 2 + offset_up[None, :],
+            mask=mask_bs[:, None],
+        )
+
+        gate = tl.where(gate > gemm1_clamp_limit, gemm1_clamp_limit, gate)
+        up = tl.where(up > gemm1_clamp_limit, gemm1_clamp_limit, up)
+        up = tl.where(up < -gemm1_clamp_limit, -gemm1_clamp_limit, up)
+        sig = 1.0 / (1.0 + tl.exp(-gate * gemm1_alpha))
+        glu = gate * sig
+        out = (up + 1) * glu
+
+        tl.store(
+            gated_output
+            + offset_bs[:, None] * output_dim
+            + tl.arange(0, output_dim)[None, :],
+            out,
+            mask=mask_bs[:, None],
+        )
+
+
+def swiglu_oai_triton(
+    hidden_states,
+    dim,
+    gemm1_alpha,
+    gemm1_clamp_limit,
+):
+    hidden_states = hidden_states.view(-1, dim)
+    BS = hidden_states.shape[0]
+    output_dim = dim // 2
+    gated_output = torch.empty(
+        (BS, output_dim),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    kernel_num = get_device_properties()[0]
+    MINIBLOCK_SIZE = 16
+    BLOCK_SIZE = triton.cdiv(BS, MINIBLOCK_SIZE * kernel_num) * MINIBLOCK_SIZE
+    BLOCK_NUM = triton.cdiv(BS, BLOCK_SIZE)
+
+    swiglu_oai_kernel[(BLOCK_NUM,)](
+        hidden_states,
+        gated_output,
+        gemm1_alpha,
+        gemm1_clamp_limit,
+        output_dim,
+        BLOCK_SIZE,
+        MINIBLOCK_SIZE,
+        BS,
+    )
+    return gated_output
+
+
+def swiglu_oai_native(layer, hidden_states):
     E, N, _ = layer.w13_weight.size()
     gate_up = hidden_states.view(-1, N)
     alpha = layer.moe_runner_config.gemm1_alpha
@@ -14,3 +93,12 @@ def swiglu_oai(layer, hidden_states):
     glu = gate * torch.sigmoid(gate * alpha)
     gated_output = (up + 1) * glu
     return gated_output
+
+
+def swiglu_oai(layer, hidden_states):
+    return swiglu_oai_triton(
+        hidden_states,
+        layer.w13_weight.shape[1],
+        layer.moe_runner_config.gemm1_alpha,
+        layer.moe_runner_config.gemm1_clamp_limit,
+    )