chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,131 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 "NvInfer.h"
+#include "common/bertCommon.h"
+#include "common/common.cuh"
+#include "common/plugin.h"
+#include "common/serialize.hpp"
+
+#include <cassert>
+#include <cstring>
+#include <cuda.h>
+#include <vector>
+
+using namespace nvinfer1;
+
+namespace nvinfer1
+{
+namespace plugin
+{
+namespace bert
+{
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelMTron(int32_t ld, int32_t const* inputIds, int32_t const* tokenIds,
+    int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb,
+    int32_t const wordSize, int32_t const tokSize, T* output, T* skip)
+{
+    // this code currently assumes the input shape is SxB, row-major => seqPos = s * B + b
+    // instead we want BxS, row-major => seqPos = b * S + s
+
+    // 1. lookup word and token of the block
+    // blockIdx.x = position in the sequence
+    // blockIdx.y = batch
+    // gridDim.x = S
+    // gridDim.y = B
+    int32_t const s = blockIdx.x;
+    int32_t const b = blockIdx.y;
+
+    int32_t const sumS = cuSeqlens[b];
+    int32_t const s_b = cuSeqlens[b + 1] - sumS;
+    if (s >= s_b)
+    {
+        return; // This CTA has nothing to do
+    }
+    __shared__ int32_t wordId;
+    __shared__ int32_t tokenId;
+
+    T const rld = T(1.f) / T(ld);
+    // seqPos = b + s * B
+    // int32_t const seqPos = blockIdx.y + blockIdx.x * gridDim.y;
+
+    // int32_t const seqPos = s * B + s;
+    int32_t const seqPos = sumS + s;
+    if (threadIdx.x == 0)
+    {
+        wordId = inputIds[seqPos];
+        tokenId = tokenIds[seqPos];
+    }
+    __syncthreads();
+
+    // 2. load pos/tok/word embeddings and add them toghether
+    // offset into embeddings is given by wordId * hidden_size
+    int32_t const poffset = blockIdx.x * ld;
+    int32_t const woffset = wordId * ld;
+    int32_t const toffset = tokenId * ld;
+    // the output offset is given by b * (S*hidden_size) + s * hidden_size
+    int32_t const outOffset = seqPos * ld;
+
+    kvp<T> threadData(0, 0);
+
+    if (wordId >= 0 && wordId < wordSize && tokenId >= 0 && tokenId < tokSize)
+    {
+        for (int32_t it = threadIdx.x; it < ld; it += TPB)
+        {
+            T const w(wordEmb[woffset + it]);
+            T const t(tokEmb[toffset + it]);
+            T const p(posEmb[poffset + it]);
+            T const val = w + t + p;
+
+            output[outOffset + it] = val;
+            skip[outOffset + it] = val;
+            T const rldval = rld * val;
+            threadData = threadData + kvp<T>(rldval, rldval * val);
+        }
+    }
+
+    // 3. layer norm on the sum
+    layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T>
+int32_t embSkipLayerNormMTron(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
+    int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
+    T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output, T* skip)
+{
+
+    constexpr int32_t tpb = 256;
+    dim3 const grid(S, B, 1);
+    dim3 const block(tpb, 1, 1);
+
+    embLayerNormKernelMTron<T, tpb><<<grid, block, 0, stream>>>(
+        ld, inputIds, tokenIds, cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, wordSize, tokSize, output, skip);
+    return cudaPeekAtLastError();
+}
+
+template int32_t embSkipLayerNormMTron<float>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
+    int32_t const*, float const*, float const*, float const*, float const*, float const*, int32_t const, int32_t const,
+    float*, float*);
+
+template int32_t embSkipLayerNormMTron<half>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
+    int32_t const*, float const*, float const*, half const*, half const*, half const*, int32_t const, int32_t const,
+    half*, half*);
+
+} // namespace bert
+} // namespace plugin
+} // namespace nvinfer1