warren/markbaseengine
1
0
Fork 0
Code Issues Pull Requests Actions 19 Packages Projects Releases Wiki Activity

v2: add embedding tests, multilingual embedding support
CI / build (push) Waiting to run
Details
CI / unit-tests (push) Blocked by required conditions
Details
CI / lint (push) Blocked by required conditions
Details

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 08:01:52 +08:00
parent ba4c41c29f
commit 85dd87e28a
13 changed files with 484 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Metal/OptimizedKernels.metal
+12 -11
View File
@@ -815,13 +815,14 @@ kernel void moe_mega_kernel(
constant uint &numExperts [[buffer(16)]],
constant float &routerScale [[buffer(17)]],
constant uint &topK [[buffer(18)]],
constant uint &groupSize [[buffer(19)]],
threadgroup float *shared_space [[threadgroup(0)]],
uint gid [[thread_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tgSize [[threads_per_threadgroup]]
) {
uint numGroupsIn = hiddenSize / 64;
uint numGroupsOut = moeIntermediate / 64;
uint numGroupsIn = hiddenSize / groupSize;
uint numGroupsOut = moeIntermediate / groupSize;
uint packedPerIn = hiddenSize / 8;
uint packedPerOut = moeIntermediate / 8;
@@ -841,10 +842,10 @@ kernel void moe_mega_kernel(
for (uint g = 0; g < numGroupsIn; g++) {
float scale = s_router[tid * numGroupsIn + g];
float bias = b_router[tid * numGroupsIn + g];
uint wBase = tid * packedPerIn + g * 8;
uint xBase = g * 64;
uint wBase = tid * packedPerIn + g * (groupSize / 8);
uint xBase = g * groupSize;
for (uint p = 0; p < 8; p += 4) {
for (uint p = 0; p < groupSize / 8; p += 4) {
device uint4 *rPtr = (device uint4*)(&w_router[wBase + p]);
uint4 packed = *rPtr;
@@ -971,10 +972,10 @@ kernel void moe_mega_kernel(
float uScale = s_up[sUpBase + gid * numGroupsIn + g];
float uBias = b_up[sUpBase + gid * numGroupsIn + g];
uint wb = gid * packedPerIn + g * 8;
uint xBase = g * 64;
uint wb = gid * packedPerIn + g * (groupSize / 8);
uint xBase = g * groupSize;
for (uint p = 0; p < 8; p += 4) {
for (uint p = 0; p < groupSize / 8; p += 4) {
device uint4 *gPtr = (device uint4*)(&w_gate[wGateBase + wb + p]);
device uint4 *uPtr = (device uint4*)(&w_up[wUpBase + wb + p]);
uint4 gP = *gPtr;
@@ -1047,10 +1048,10 @@ kernel void moe_mega_kernel(
float scale = s_down[wDownBase + gid * numGroupsOut + g];
float bias = b_down[wDownBase + gid * numGroupsOut + g];
uint wb = gid * packedPerOut + g * 8;
uint iBase = g * 64;
uint wb = gid * packedPerOut + g * (groupSize / 8);
uint iBase = g * groupSize;
for (uint p = 0; p < 8; p += 4) {
for (uint p = 0; p < groupSize / 8; p += 4) {
device uint4 *wPtr = (device uint4*)(&w_down[wDownBase + wb + p]);
uint4 packed = *wPtr;