const int64_t row_idx = blockIdx.x;

const int64_t tid = threadIdx.x;

const BFloat16* o1_row = o1 + row_idx * (int64_t)moe_intermediate_size * 2;

const BFloat16* do2_s_row = do2_s + row_idx * (int64_t)moe_intermediate_size;

reinterpret_cast<const bfloat16x4_t*>(o1_row +

(int64_t)moe_intermediate_size);

BFloat16* do1_row = do1 + row_idx * (int64_t)moe_intermediate_size * 2;

BFloat16* o2s_row = o2_s + row_idx * (int64_t)moe_intermediate_size;

const int vec_numel = (int64_t)moe_intermediate_size / numel_per_thread;

for (int64_t i = tid; i < vec_numel; i += blockDim.x) {