Add chunk kda fwd

build.sh

@@ -399,13 +399,24 @@ function build_cmake_modules()
         )
     fi
 
-    rm -rf "$BUILD_DIR"
+    # KEEP_BUILD_DIR=1 reuses the existing CMake cache and object files. The
+    # AscendC kernel targets (chunk_kda_fwd in particular) dominate a cold build,
+    # so keeping them lets an unrelated edit rebuild in seconds. Default stays a
+    # clean build; drop the variable whenever the CMake configuration changes.
+    if [[ "${KEEP_BUILD_DIR:-0}" == "1" && -f "$BUILD_DIR/CMakeCache.txt" ]]; then
+        echo "KEEP_BUILD_DIR=1: reusing $BUILD_DIR (incremental build)"
+    else
+        rm -rf "$BUILD_DIR"
+    fi
     mkdir -p "$BUILD_DIR"
 
     cmake "${cmake_args[@]}" \
         -B "$BUILD_DIR" \
         -S "$PROJECT_ROOT"
-    cmake --build "$BUILD_DIR" --target install -j 16
+    # Parallelism is bounded by RAM, not cores: the CCE compiler needs ~1-2 GB per
+    # job and the AscendC kernel TUs are the memory-hungry ones. Keep the historical
+    # default and let callers tune it with BUILD_JOBS.
+    cmake --build "$BUILD_DIR" --target install -j "${BUILD_JOBS:-16}"
 }
 
 function build_deepep_kernels()

csrc/CMakeLists.txt

@@ -57,6 +57,7 @@ set(OP_SRCS
     ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv_metadata/op_host/sparse_attn_sharedkv_metadata.cpp
     ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/sparse_attention_score.cpp
     ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/tiling/sparse_attention_score_tiling.cpp
+    ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp
     )
 # A3-only host sources
 if(SGL_KERNEL_ENABLE_A3_ONLY_OPS)
@@ -189,6 +190,32 @@ ascendc_compile_definitions(workspace_kernel PRIVATE
     -DHAVE_TILING
 )
 
+# chunk_kda_fwd: dedicated kernel target with its own compile options. The kernel
+# manages cross-pipe synchronization manually (SetFlag/WaitFlag/SyncAll), so the
+# aclnn build compiles it with --cce-auto-sync=off; keep that here instead of
+# applying it to the shared workspace_kernel target.
+ascendc_library(chunk_kda_fwd_kernel STATIC
+    ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp
+)
+ascendc_include_directories(chunk_kda_fwd_kernel PRIVATE
+    ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel
+    ${PROJECT_OP_SRC_BASE}/utils/kernel
+    ${ASCEND_INCLUDE_DIR}
+)
+ascendc_compile_definitions(chunk_kda_fwd_kernel PRIVATE
+    -DHAVE_WORKSPACE
+    -DHAVE_TILING
+)
+ascendc_compile_options(chunk_kda_fwd_kernel PRIVATE
+    --cce-auto-sync=off
+    -Wno-deprecated-declarations
+)
+if(_sgl_soc_lower MATCHES "^ascend950")
+    ascendc_compile_definitions(chunk_kda_fwd_kernel PRIVATE
+        -DENABLE_CV_COMM_VIA_SSBUF=true
+    )
+endif()
+
 # Build catlass kernel in its own target (no sparse_attention_score include dir),
 # so catlass's #include "tla/layout.hpp" resolves to its own copy and avoids
 # symbol redefinition (RoundDown/CeilDiv/Max/Min/...) with sparse_attention_score.
@@ -258,6 +285,7 @@ target_link_libraries(${OP_PLUGIN_NAME} PRIVATE
         ${A3_ONLY_KERNEL_TARGETS}
         workspace_kernel
         no_workspace_kernel
+        chunk_kda_fwd_kernel
         ${CATLASS_WORKSPACE_KERNEL_TARGETS}
         torch_npu
         ascendcl
@@ -289,6 +317,8 @@ target_include_directories(${OP_PLUGIN_NAME} PRIVATE
         ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv/op_host
         ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host
         ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/tiling
+        ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_host
+        ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel
         ${TORCH_DIR}/include
         ${TORCH_DIR}/include/torch/csrc/api/include
         ${TORCH_NPU_DIR}/include

csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h

@@ -0,0 +1,53 @@
+/**
+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#pragma once
+
+namespace optiling::arch35 {
+
+struct ChunkKdaFwdArch35Options {
+    bool computeGateInPrepare = false;
+    bool fusePostWu = false;
+    bool fusePostWuIntoFwdH = false;
+    bool useDenseFwdH = false;
+};
+
+inline ChunkKdaFwdArch35Options ConfigureChunkKdaFwdArch35(
+    bool isAscend950, bool qIsBf16, bool rawGIsFp32, bool hasALog,
+    bool useGateInKernel, bool safeGate, bool isVarLen, int64_t seqlen,
+    int64_t vHeads, int64_t chunkSize, int64_t kDim, int64_t vDim,
+    bool storeQG, bool storeVNew, bool storeH)
+{
+    ChunkKdaFwdArch35Options options;
+    const bool shapeSupported =
+        isAscend950 && chunkSize == 64 && kDim == 128 && vDim == 128;
+    if (!shapeSupported) {
+        return options;
+    }
+
+    // Tiling keys describe shape families independently of the SoC. These
+    // options only enable arch35 sub-pipelines within the selected family.
+    options.computeGateInPrepare =
+        qIsBf16 && rawGIsFp32 && hasALog &&
+        useGateInKernel && safeGate;
+    const bool denseAligned = !isVarLen && seqlen % chunkSize == 0;
+    options.useDenseFwdH = denseAligned && qIsBf16;
+    const bool canFusePreparePostWu =
+        denseAligned && qIsBf16 && safeGate && vHeads % 2 == 0;
+    options.fusePostWuIntoFwdH =
+        options.useDenseFwdH && canFusePreparePostWu &&
+        options.computeGateInPrepare &&
+        !storeQG && !storeVNew && !storeH;
+    options.fusePostWu =
+        canFusePreparePostWu && !options.fusePostWuIntoFwdH;
+    return options;
+}
+
+} // namespace optiling::arch35