mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
38 lines
1.6 KiB
C++
38 lines
1.6 KiB
C++
#include <hip/hip_runtime.h>
|
|
#include <hip/hip_bf16.h>
|
|
|
|
// GPT-OSS production keeps two 8,192-token sequences on each DP rank. One wave
|
|
// owns one token row, so its 16-byte weight reads and output writes are fully
|
|
// contiguous instead of interleaving four unrelated vocabulary rows per wave.
|
|
using uint4v = unsigned int __attribute__((ext_vector_type(4)));
|
|
|
|
constexpr int TOKENS = 16384;
|
|
constexpr int VOCAB = 128256;
|
|
constexpr int EMBED = 2880;
|
|
constexpr int VECS_PER_ROW = EMBED / 8;
|
|
constexpr int THREADS = 512;
|
|
constexpr int WAVE = 64;
|
|
constexpr int ROWS_PER_WG = THREADS / WAVE;
|
|
|
|
static_assert(EMBED % 8 == 0 && TOKENS % ROWS_PER_WG == 0);
|
|
|
|
extern "C" __global__ __launch_bounds__(THREADS) void gptoss_embedding_fwd(
|
|
__hip_bfloat16 *__restrict__ out, const int *__restrict__ idx,
|
|
const __hip_bfloat16 *__restrict__ weight) {
|
|
const int lid = threadIdx.x;
|
|
const int lane = lid & (WAVE - 1);
|
|
const int row = blockIdx.x * ROWS_PER_WG + (lid / WAVE);
|
|
const int token = idx[row];
|
|
const uint4v zero = {0u, 0u, 0u, 0u};
|
|
|
|
#pragma unroll
|
|
for (int j = lane; j < VECS_PER_ROW; j += WAVE) {
|
|
// The embedding is a one-use, random-row source. Bypass temporal caching so
|
|
// it does not displace the freshly written activation consumed by RMSNorm.
|
|
const uint4v value = static_cast<unsigned>(token) < VOCAB ?
|
|
__builtin_nontemporal_load(reinterpret_cast<const uint4v *>(weight) +
|
|
static_cast<long long>(token) * VECS_PER_ROW + j) : zero;
|
|
reinterpret_cast<uint4v *>(out)[static_cast<long long>(row) * VECS_PER_ROW + j] = value;
|
|
}
|
|
}
|