Add AscendC triangular inverse

csrc/CMakeLists.txt

@@ -21,6 +21,7 @@ FILE(GLOB OP_SRCS
     ${PROJECT_OP_SRC_BASE}/lora/op_host/sgemmv_shrink.cpp
     ${PROJECT_OP_SRC_BASE}/lightning_indexer/op_host/lightning_indexer.cpp
     ${PROJECT_OP_SRC_BASE}/lightning_indexer/op_host/tiling/lightning_indexer_tiling.cpp
+    ${PROJECT_OP_SRC_BASE}/tri_inv/op_host/tri_inv.cpp
     )
 if(BUILD_CATLASS_MODULE)
     list(APPEND OP_SRCS
@@ -43,6 +44,7 @@ ascendc_library(no_workspace_kernel STATIC
     ${PROJECT_OP_SRC_BASE}/lora/op_kernel/sgmv_shrink_kernel.cpp
     ${PROJECT_OP_SRC_BASE}/lora/op_kernel/sgemmv_expand_kernel.cpp
     ${PROJECT_OP_SRC_BASE}/lora/op_kernel/sgemmv_shrink_kernel.cpp
+    ${PROJECT_OP_SRC_BASE}/tri_inv/op_kernel/tri_inv_kernel.cpp
 )
 
 # kernel side files with workspace

csrc/pytorch_extensions.cpp

@@ -96,6 +96,8 @@ TORCH_LIBRARY_FRAGMENT(npu, m)
         "Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, "
         "str? layout_query=None, str? layout_key=None, "
         "int? sparse_count=None, int? sparse_mode=None) -> Tensor");
+
+    m.def("triangular_inverse(Tensor x) -> Tensor");
 }
 }  // namespace
 
@@ -137,5 +139,7 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m)
 #endif
 
     m.impl("lightning_indexer", TORCH_FN(sglang::npu_kernel::lightning_indexer));
+
+    m.impl("triangular_inverse", TORCH_FN(sglang::npu_kernel::tri_inv_col_sweep));
 }
 }  // namespace

csrc/tri_inv/README.md

@@ -0,0 +1,5 @@
+##### Description of tri_inv
+
+This is a vector-only AscendC triangular inversion kernel on Ascend NPU.
+
+The kernel supports matrix sizes `16, 32, 64, 128` and data types `fp16` and `fp32`.

csrc/tri_inv/op_host/tiling_tri_inv.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <cstdint>
+
+namespace sglang {
+
+namespace npu_kernel {
+
+/**
+ * @brief `tri_inv_col_sweep` kernel tiling parameter structure.
+ */
+struct TriInvColumnSweepTiling {
+    /// @brief Number of blocks.
+    uint32_t num_blocks;
+    /// @brief Total number of input elements.
+    uint32_t num_elems;
+    /// @brief Input matrix size.
+    uint32_t matrix_size;
+};
+
+}  // namespace npu_kernel
+}  // namespace sglang

csrc/tri_inv/op_host/tri_inv.cpp

@@ -0,0 +1,71 @@
+// Licensed under the BSD 3-Clause License  (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "defines.h"
+#include "torch_helper.h"
+
+#include "tiling_tri_inv.h"
+#include "aclrtlaunch_tri_inv_col_sweep_fp16.h"
+#include "aclrtlaunch_tri_inv_col_sweep_fp32.h"
+
+namespace sglang {
+
+namespace npu_kernel {
+
+at::Tensor calc_tiling(const TriInvColumnSweepTiling &tiling)
+{
+    constexpr uint32_t PADDING_BYTE = 32U;
+
+    // align to 32 bytes
+    int32_t tiling_size = (sizeof(TriInvColumnSweepTiling) + PADDING_BYTE - 1) / PADDING_BYTE * PADDING_BYTE;
+    auto tiling_buffer = at::empty({tiling_size}, at::TensorOptions().dtype(at::kByte).device(at::kCPU));
+
+    TriInvColumnSweepTiling *tiling_data = reinterpret_cast<TriInvColumnSweepTiling *>(tiling_buffer.data_ptr());
+    tiling_data->num_blocks = tiling.num_blocks;
+    tiling_data->num_elems = tiling.num_elems;
+    tiling_data->matrix_size = tiling.matrix_size;
+
+    auto tiling_tensor = TorchNpuHelper::CopyTensorHostToDevice(tiling_buffer);
+    return tiling_tensor;
+}
+
+HOST_API at::Tensor tri_inv_col_sweep(const at::Tensor &tensor)
+{
+    const auto dtype = tensor.options().dtype();
+    if (tensor.dim() < 2) {
+        throw std::runtime_error("Input tensor must have at least 2 dimensions.\n");
+    }
+
+    const uint32_t matrix_size = static_cast<uint32_t>(tensor.size(-1));
+    if (matrix_size != tensor.size(-2)) {
+        throw std::runtime_error("Only square matrices are supported.\n");
+    }
+
+    const uint32_t num_elems = static_cast<uint32_t>(tensor.numel());
+    const uint32_t block_dim = static_cast<uint32_t>(num_elems / (matrix_size * matrix_size));
+
+    const at::Tensor tensor_out = at::empty_like(tensor);
+
+    const TriInvColumnSweepTiling tiling{block_dim, num_elems, matrix_size};
+    const at::Tensor tiling_device = calc_tiling(tiling);
+
+    if (dtype == at::kHalf) {
+        EXEC_KERNEL_CMD(tri_inv_col_sweep_fp16, block_dim, tensor, tensor_out, tiling_device);
+    } else if (dtype == at::kFloat) {
+        EXEC_KERNEL_CMD(tri_inv_col_sweep_fp32, block_dim, tensor, tensor_out, tiling_device);
+    } else {
+        throw std::runtime_error("Unsupported data type for tri_inv_col_sweep. fp16 and fp32 are currently supported.");
+    }
+
+    return tensor_out;
+}
+
+}  // namespace npu_kernel
+}  // namespace sglang

csrc/tri_inv/op_kernel/kernel_tri_inv.h

@@ -0,0 +1,184 @@
+// Licensed under the BSD 3-Clause License  (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
+ *
+ * @file kernel_tri_inv.h
+ * @brief Kernel implementing a Vector matrix inverse kernel operation.
+ */
+
+#pragma once
+
+#include "kernel_operator.h"
+
+namespace sglang {
+
+namespace npu_kernel {
+/**
+ * @brief Returns the matrix inverse of an upper triangular square matrix of
+ * size `matrix_size`. The matrix has ones on the main diagonal.
+ *
+ * The column sweep algorithm is used for the linear system Ax=e_j where e_j is
+ * the standard vector.
+ *
+ * @tparam T Input data type. Supports only `half` and `float32`.
+ *
+ */
+template <typename T>
+class KernelTriInvColumnSweep
+{
+    constexpr static uint32_t BUFFER_NUM = 1;
+
+public:
+    /**
+     * @brief Class constructor.
+     *
+     * @param [in] vec_len Total length of input tensor.
+     * @param [in] matrix_size Input square matrix size.
+     */
+    __aicore__ inline KernelTriInvColumnSweep(uint32_t vec_len, uint32_t matrix_size)
+        : vec_len_(vec_len), matrix_size_(matrix_size), tile_len_(matrix_size * matrix_size)
+    {}
+
+    /**
+     * @brief Initialize global and local memory structures.
+     *
+     * @param [in] vec_in Pointer to the input vector in global memory.
+     * @param [in] vec_out Pointer to the output vector in global memory.
+     */
+    __aicore__ inline void Init(GM_ADDR vec_in, GM_ADDR vec_out)
+    {
+        global_in_.SetGlobalBuffer((__gm__ T *)vec_in, vec_len_);
+        global_out_.SetGlobalBuffer((__gm__ T *)vec_out, vec_len_);
+
+        pipe_.InitBuffer(in_q_, BUFFER_NUM, tile_len_ * sizeof(T));
+        pipe_.InitBuffer(out_q_, BUFFER_NUM, tile_len_ * sizeof(T));
+        pipe_.InitBuffer(b_buf_, matrix_size_ * sizeof(T));
+    }
+
+    /**
+     * @brief Run the kernel.
+     */
+    __aicore__ inline void Process()
+    {
+        using namespace AscendC;
+        const uint32_t global_offset = AscendC::GetBlockIdx() * tile_len_;
+
+        const AscendC::LocalTensor<T> tile_in_lt = in_q_.AllocTensor<T>();
+        AscendC::DataCopy(tile_in_lt, global_in_[global_offset], tile_len_);
+        in_q_.EnQue(tile_in_lt);
+
+        InvertMatrix();
+
+        AscendC::LocalTensor<T> tile_out_lt = out_q_.DeQue<T>();
+        AscendC::DataCopy(global_out_[global_offset], tile_out_lt, tile_len_);
+        out_q_.FreeTensor(tile_out_lt);
+    }
+
+private:
+    __aicore__ inline void InvertMatrix()
+    {
+        using namespace AscendC;
+
+        const int32_t n_rows = matrix_size_;
+        const int32_t n_cols = matrix_size_;
+
+        LocalTensor<T> vec_in_lt = in_q_.DeQue<T>();
+        const LocalTensor<T> vec_out_lt = out_q_.AllocTensor<T>();
+
+        // Left-hand side Ax=b.
+        LocalTensor<T> b = b_buf_.Get<T>();
+
+        Duplicate(vec_out_lt, static_cast<T>(0), tile_len_);
+
+        // For every output column j-th
+        for (int32_t j = 0; j < n_cols; j++) {
+            // Column sweep on each column.
+
+            // `b` vector is e_j standard vector.
+            Duplicate(b, static_cast<T>(0), matrix_size_);
+            b.SetValue(j, static_cast<T>(1));
+
+            // Ax=b
+            LocalTensor<T> x = vec_out_lt[j * n_rows];
+            for (int32_t k = n_rows - 1; k >= 0; k--) {
+                const LocalTensor<T> A_k = vec_in_lt[k * n_rows];
+
+                // x[k] = b[k] / A[k, k]
+                x.SetValue(k, b.GetValue(k));
+
+                if (k > 0) {
+                    // b[:k] -= A[:k, k] * x[k]
+                    const float x_k = -static_cast<float>(x.GetValue(k));
+                    AscendC::Axpy<T>(b, A_k, static_cast<T>(x_k), k);
+                }
+            }
+        }
+
+        out_q_.EnQue<T>(vec_out_lt);
+        in_q_.FreeTensor<T>(vec_in_lt);
+    }
+
+    AscendC::TPipe pipe_;
+
+    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> in_q_;
+    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> out_q_;
+
+    AscendC::TBuf<AscendC::QuePosition::VECCALC> b_buf_;
+
+    AscendC::GlobalTensor<T> global_in_;
+    AscendC::GlobalTensor<T> global_out_;
+
+    const uint32_t vec_len_;
+    const uint32_t matrix_size_;
+    const uint32_t tile_len_;
+};
+
+/**
+ * @brief Run the `tri_inv_col_sweep` kernel.
+ *
+ * @tparam T Input data type. Supports fp16/half.
+ *
+ * @param [in] vec_in Pointer to the input vector.
+ * @param [in] vec_out Pointer ot the output vector.
+ * @param [in] vec_len Dimension of the input vector.
+ * @param [in] matrix_size Matrix size to invert.
+ */
+template <typename T>
+__aicore__ inline void run_tri_inv_col_sweep(GM_ADDR vec_in, GM_ADDR vec_out, uint32_t vec_len, uint32_t matrix_size)
+{
+    if ASCEND_IS_AIV {
+        KernelTriInvColumnSweep<T> op(vec_len, matrix_size);
+        op.Init(vec_in, vec_out);
+        op.Process();
+    }
+}
+
+/**
+ * @brief Copies tiling structure from global memory to registers.
+ *
+ * @tparam TilingT Structure representing kernel tiling parameters.
+ * @param [in] tiling Pointer to the structure allocated in registers.
+ * @param [in] tiling_global Pointer to the structure in global memory.
+ */
+template <typename TilingT>
+__aicore__ inline void GetTilingData(TilingT *const tiling, GM_ADDR tiling_global)
+{
+    uint32_t *const tiling_32b = reinterpret_cast<uint32_t *>(tiling);
+    const __gm__ uint32_t *const tiling_global_32b = reinterpret_cast<__gm__ uint32_t *>(tiling_global);
+
+    for (uint32_t i = 0; i < sizeof(TilingT) / sizeof(uint32_t); i++) {
+        tiling_32b[i] = tiling_global_32b[i];
+    }
+}
+
+}  // namespace npu_kernel
+}  // namespace sglang

csrc/tri_inv/op_kernel/tri_inv_kernel.cpp

@@ -0,0 +1,41 @@
+// Licensed under the BSD 3-Clause License  (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "kernel_tri_inv.h"
+
+#include "../op_host/tiling_tri_inv.h"
+
+/**
+ * @brief Run the `tri_inv_col_sweep` kernel on dtype fp16/half.
+ *
+ * @param [in] vec_in Pointer to input vector.
+ * @param [in] vec_out Pointer to output vector.
+ * @param [in] tiling_gm Pointer to tiling vector.
+ */
+extern "C" __global__ __aicore__ void tri_inv_col_sweep_fp16(GM_ADDR vec_in, GM_ADDR vec_out, GM_ADDR tiling_gm)
+{
+    sglang::npu_kernel::TriInvColumnSweepTiling tiling;
+    sglang::npu_kernel::GetTilingData(&tiling, tiling_gm);
+    sglang::npu_kernel::run_tri_inv_col_sweep<half>(vec_in, vec_out, tiling.num_elems, tiling.matrix_size);
+}
+
+/**
+ * @brief Run the `tri_inv_col_sweep` kernel on dtype float32.
+ *
+ * @param [in] vec_in Pointer to input vector.
+ * @param [in] vec_out Pointer to output vector.
+ * @param [in] tiling_gm Pointer to tiling vector.
+ */
+extern "C" __global__ __aicore__ void tri_inv_col_sweep_fp32(GM_ADDR vec_in, GM_ADDR vec_out, GM_ADDR tiling_gm)
+{
+    sglang::npu_kernel::TriInvColumnSweepTiling tiling;
+    sglang::npu_kernel::GetTilingData(&tiling, tiling_gm);
+    sglang::npu_kernel::run_tri_inv_col_sweep<float>(vec_in, vec_out, tiling.num_elems, tiling.matrix_size);
+}

include/sgl_kenel_npu_ops.h

@@ -113,6 +113,16 @@ at::Tensor lightning_indexer(
     c10::optional<c10::string_view> layout_key,
     c10::optional<int64_t> sparse_count, c10::optional<int64_t> sparse_mode);
 
+/**
+ * @brief Triangular inverse of input tensor where last two dimensions represent
+ * a matrix.
+ *
+ * @param [in] tensor_in Tensor of dimensions (..., n, n) where `n` is
+ * the matrix size.
+ * @return at::Tensor Returns tensor of same shape where each matrix of size n
+ * is inversed.
+ */
+at::Tensor tri_inv_col_sweep(const at::Tensor &tensor_in);
 } // namespace npu_kernel
 
 } // namespace sglang

python/sgl_kernel_npu/setup.py

@@ -13,8 +13,6 @@
 from setuptools.dist import Distribution
 from torch_npu.utils.cpp_extension import NpuExtension
 
-os.environ["SOURCE_DATE_EPOCH"] = "0"
-
 
 class BinaryDistribution(Distribution):
     """Distribution which always forces a binary package with platform name"""