| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/mhc.cpp |
| 3 | // |
| 4 | // DeepSeek-V4 mHC (multi-residual-stream hyper-connection) gating kernels. |
| 5 | // |
| 6 | // Difference from upstream: vLLM's `mhc_pre_torch` reference (kernels/ |
| 7 | // mhc/torch.py) exposes separate `hc_pre_eps`/`hc_sinkhorn_eps`/ |
| 8 | // `hc_post_mult_value`/`n_splits` knobs this kernel doesn't implement -- |
| 9 | // only ever called with hc_pre_eps == hc_sinkhorn_eps, hc_post_mult_value |
| 10 | // == 2.0, n_splits == 1 (matches upstream), so the entry points take one |
| 11 | // merged `hc_eps` and hardcode the post-mix multiplier to 2.0. The Python |
| 12 | // wrapper (kernels/mhc/cpu.py) enforces the merged-eps assumption via |
| 13 | // TORCH_CHECK. |
| 14 | // |
| 15 | // clang-format off |
| 16 | |
| 17 | #include <algorithm> |
| 18 | #include <cmath> |
| 19 | #include <limits> |
| 20 | |
| 21 | #include "common.h" |
| 22 | #include "vec.h" |
| 23 | |
| 24 | namespace { |
| 25 | |
| 26 | // Loop tiling heuristics. |
| 27 | constexpr int64_t kMhcTileCacheBudgetBytes = 36 * 1024; |
| 28 | constexpr int64_t kMhcDefaultTokenBlock = 32; |
| 29 | |
| 30 | inline int64_t round_to_nearest_multiple(int64_t value, int64_t multiple) { |
| 31 | return ((value + multiple / 2) / multiple) * multiple; |
| 32 | } |
| 33 | |
| 34 | inline int64_t choose_t_block(int64_t /*T*/) { |
| 35 | return kMhcDefaultTokenBlock; |
| 36 | } |
| 37 | |
| 38 | inline int64_t choose_k_block(int64_t weight_rows, int64_t vec_size, int64_t max_block) { |
| 39 | const int64_t bytes_per_k = (weight_rows + 1) * static_cast<int64_t>(sizeof(float)); |
| 40 | const int64_t raw = std::max(vec_size, kMhcTileCacheBudgetBytes / bytes_per_k); |
| 41 | const int64_t rounded = round_to_nearest_multiple(raw, vec_size); |
| 42 | return std::max(vec_size, std::min(max_block, rounded)); |
| 43 | } |
| 44 | |
| 45 | template <typename func_t> |
| 46 | inline void parallel_mhc_phase1(int64_t n_t_blocks, int64_t n_k_blocks, const func_t& f) { |
| 47 | const int64_t nth = static_cast<int64_t>(at::get_num_threads()); |
| 48 | if (n_t_blocks >= nth) { |
| 49 | // Large prefill: enough token blocks exist to fill all workers. |
| 50 | parallel_2d_tiled(static_cast<int>(n_t_blocks), static_cast<int>(n_k_blocks), static_cast<int>(nth), 1, f); |
| 51 | } else { |
| 52 | // Decode / small T: split both T and K to expose enough parallel tasks. |
| 53 | parallel_2d(static_cast<int>(n_t_blocks), static_cast<int>(n_k_blocks), f); |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | // In-place Sinkhorn on a row-major [hc, hc] matrix. |
| 58 | inline void |
| 59 | sinkhorn_inplace(float* __restrict__ comb_matrix, int hc, int sinkhorn_iters, float eps, float* __restrict__ scratch) { |
| 60 | float* const row_sum = scratch; |
| 61 | float* const col_sum = scratch + hc; |
| 62 | |
| 63 | // First iteration. |
| 64 | for (int r = 0; r < hc; ++r) { |
| 65 | const float* row = comb_matrix + r * hc; |
| 66 | float row_acc = 0.f; |
| 67 | #pragma GCC unroll 4 |
| 68 | for (int c = 0; c < hc; ++c) |
| 69 | row_acc += row[c]; |
| 70 | row_sum[r] = row_acc; |
| 71 | } |
| 72 | for (int r = 0; r < hc; ++r) { |
| 73 | const float row_inv = 1.0f / row_sum[r]; |
| 74 | float* row = comb_matrix + r * hc; |
| 75 | #pragma GCC unroll 4 |
| 76 | for (int c = 0; c < hc; ++c) |
| 77 | row[c] = row[c] * row_inv + eps; |
| 78 | } |
| 79 | for (int c = 0; c < hc; ++c) { |
| 80 | float col_acc = 0.f; |
| 81 | #pragma GCC unroll 4 |
| 82 | for (int r = 0; r < hc; ++r) |
| 83 | col_acc += comb_matrix[r * hc + c]; |
| 84 | col_sum[c] = col_acc; |
| 85 | } |
| 86 | for (int r = 0; r < hc; ++r) |
| 87 | #pragma GCC unroll 4 |
| 88 | for (int c = 0; c < hc; ++c) |
| 89 | comb_matrix[r * hc + c] /= (col_sum[c] + eps); |
| 90 | |
| 91 | // Remaining iterations. |
| 92 | for (int iter = 1; iter < sinkhorn_iters; ++iter) { |
| 93 | for (int r = 0; r < hc; ++r) { |
| 94 | const float* row = comb_matrix + r * hc; |
| 95 | float row_acc = 0.f; |
| 96 | #pragma GCC unroll 4 |
| 97 | for (int c = 0; c < hc; ++c) |
| 98 | row_acc += row[c]; |
| 99 | row_sum[r] = row_acc; |
| 100 | } |
| 101 | for (int r = 0; r < hc; ++r) { |
| 102 | const float row_inv = 1.0f / (row_sum[r] + eps); |
| 103 | float* row = comb_matrix + r * hc; |
| 104 | #pragma GCC unroll 4 |
| 105 | for (int c = 0; c < hc; ++c) |
| 106 | row[c] *= row_inv; |
| 107 | } |
| 108 | for (int c = 0; c < hc; ++c) { |
| 109 | float col_acc = 0.f; |
| 110 | #pragma GCC unroll 4 |
| 111 | for (int r = 0; r < hc; ++r) |
| 112 | col_acc += comb_matrix[r * hc + c]; |
| 113 | col_sum[c] = col_acc; |
| 114 | } |
| 115 | for (int r = 0; r < hc; ++r) |
| 116 | #pragma GCC unroll 4 |
| 117 | for (int c = 0; c < hc; ++c) |
| 118 | comb_matrix[r * hc + c] /= (col_sum[c] + eps); |
| 119 | } |
| 120 | } |
| 121 | |
| 122 | // Parse mixes -> pre/post/comb and run sinkhorn on comb. |
| 123 | inline void parse_mixes_and_sinkhorn( |
| 124 | float* __restrict__ pre_gate, // [hc] output |
| 125 | float* __restrict__ post_gate, // [hc] output |
| 126 | float* __restrict__ comb_matrix, // [hc*hc] output (row-major) |
| 127 | const float* __restrict__ mixes, // [mix_hc] input |
| 128 | float s0, |
| 129 | float s1, |
| 130 | float s2, |
| 131 | const float* __restrict__ hc_base, |
| 132 | int sinkhorn_iters, |
| 133 | float eps, |
| 134 | int hc, |
| 135 | float* __restrict__ scratch) { |
| 136 | #pragma GCC unroll 4 |
| 137 | for (int h = 0; h < hc; ++h) { |
| 138 | const float v = mixes[h] * s0 + hc_base[h]; |
| 139 | pre_gate[h] = 1.f / (1.f + std::exp(-v)) + eps; |
| 140 | } |
| 141 | #pragma GCC unroll 4 |
| 142 | for (int h = 0; h < hc; ++h) { |
| 143 | const float v = mixes[hc + h] * s1 + hc_base[hc + h]; |
| 144 | post_gate[h] = 2.f / (1.f + std::exp(-v)); |
| 145 | } |
| 146 | float* const row_max = scratch; |
| 147 | for (int r = 0; r < hc; ++r) |
| 148 | row_max[r] = -std::numeric_limits<float>::infinity(); |
| 149 | for (int r = 0; r < hc; ++r) { |
| 150 | float* row = comb_matrix + r * hc; |
| 151 | for (int c = 0; c < hc; ++c) { |
| 152 | const int idx = r * hc + c; |
| 153 | const float v = mixes[2 * hc + idx] * s2 + hc_base[2 * hc + idx]; |
| 154 | row[c] = v; |
| 155 | if (v > row_max[r]) row_max[r] = v; |
| 156 | } |
| 157 | } |
| 158 | for (int r = 0; r < hc; ++r) { |
| 159 | float* row = comb_matrix + r * hc; |
| 160 | #pragma GCC unroll 4 |
| 161 | for (int c = 0; c < hc; ++c) |
| 162 | row[c] = std::exp(row[c] - row_max[r]); |
| 163 | } |
| 164 | |
| 165 | sinkhorn_inplace(comb_matrix, hc, sinkhorn_iters, eps, scratch + hc); |
| 166 | } |
| 167 | |
| 168 | template <typename scalar_t> |
| 169 | static void hc_post_fuse_impl( |
| 170 | scalar_t* __restrict__ out, // [T, hc, d] output |
| 171 | const scalar_t* __restrict__ x, // [T, d] input |
| 172 | const scalar_t* __restrict__ residual, // [T, hc, d] input |
| 173 | const float* __restrict__ post, // [T, hc] float32 |
| 174 | const float* __restrict__ comb, // [T, hc, hc] float32 row-major |
| 175 | int64_t T, |
| 176 | int64_t d, |
| 177 | int hc) { |
| 178 | using bVec = at::vec::Vectorized<scalar_t>; |
| 179 | using fVec = at::vec::Vectorized<float>; |
| 180 | constexpr int64_t kVecSize = bVec::size(); |
| 181 | |
| 182 | const int64_t nthreads = at::get_num_threads(); |
| 183 | const int64_t max_splits = std::max(int64_t(1), d / kVecSize); |
| 184 | const int64_t token_head_tasks = T * hc; |
| 185 | const int64_t k_splits = (token_head_tasks >= nthreads) ? int64_t(1) : std::min(nthreads, max_splits); |
| 186 | const int64_t k_chunk = (d + k_splits - 1) / k_splits; |
| 187 | |
| 188 | const int64_t fvec_floats = static_cast<int64_t>(sizeof(fVec) / sizeof(float)); |
| 189 | auto comb_fvec_storage = at::empty({nthreads * hc * fvec_floats}, at::kFloat); |
| 190 | fVec* const comb_fvec_base = reinterpret_cast<fVec*>(comb_fvec_storage.data_ptr<float>()); |
| 191 | |
| 192 | at::parallel_for(0, T * hc * k_splits, 0, [&](int64_t begin, int64_t end) { |
| 193 | const int tid = static_cast<int>(at::get_thread_num()); |
| 194 | fVec* const comb_fvec = comb_fvec_base + tid * hc; |
| 195 | |
| 196 | for (int64_t idx = begin; idx < end; ++idx) { |
| 197 | const int64_t task = idx / k_splits; |
| 198 | const int64_t split_id = idx % k_splits; |
| 199 | const int64_t t = task / hc; |
| 200 | const int h = static_cast<int>(task % hc); |
| 201 | const int64_t k0 = split_id * k_chunk; |
| 202 | const int64_t k1 = std::min(k0 + k_chunk, d); |
| 203 | |
| 204 | const float post_val = post[t * hc + h]; |
| 205 | const float* comb_t = comb + t * hc * hc; |
| 206 | const scalar_t* x_t = x + t * d; |
| 207 | const scalar_t* res_t = residual + t * hc * d; |
| 208 | scalar_t* out_th = out + (t * hc + h) * d; |
| 209 | |
| 210 | const fVec post_fvec(post_val); |
| 211 | #pragma GCC unroll 4 |
| 212 | for (int i = 0; i < hc; ++i) |
| 213 | comb_fvec[i] = fVec(comb_t[i * hc + h]); |
| 214 | |
| 215 | int64_t k; |
| 216 | // bf16: convert and accumulate |
| 217 | for (k = k0; k <= k1 - kVecSize; k += kVecSize) { |
| 218 | fVec acc0, acc1; |
| 219 | std::tie(acc0, acc1) = at::vec::convert_to_float(bVec::loadu(x_t + k)); |
| 220 | acc0 = post_fvec * acc0; |
| 221 | acc1 = post_fvec * acc1; |
| 222 | #pragma GCC unroll 4 |
| 223 | for (int i = 0; i < hc; ++i) { |
| 224 | fVec r0, r1; |
| 225 | std::tie(r0, r1) = at::vec::convert_to_float(bVec::loadu(res_t + i * d + k)); |
| 226 | acc0 += comb_fvec[i] * r0; |
| 227 | acc1 += comb_fvec[i] * r1; |
| 228 | } |
| 229 | at::vec::convert_from_float<scalar_t>(acc0, acc1).store(out_th + k); |
| 230 | } |
| 231 | if (k < k1) { |
| 232 | const int64_t rem = k1 - k; |
| 233 | fVec acc0, acc1; |
| 234 | std::tie(acc0, acc1) = at::vec::convert_to_float(bVec::loadu(x_t + k, rem)); |
| 235 | acc0 = post_fvec * acc0; |
| 236 | acc1 = post_fvec * acc1; |
| 237 | #pragma GCC unroll 4 |
| 238 | for (int i = 0; i < hc; ++i) { |
| 239 | fVec r0, r1; |
| 240 | std::tie(r0, r1) = at::vec::convert_to_float(bVec::loadu(res_t + i * d + k, rem)); |
| 241 | acc0 += comb_fvec[i] * r0; |
| 242 | acc1 += comb_fvec[i] * r1; |
| 243 | } |
| 244 | at::vec::convert_from_float<scalar_t>(acc0, acc1).store(out_th + k, rem); |
| 245 | } |
| 246 | } |
| 247 | }); |
| 248 | } |
| 249 | |
| 250 | // Fused hc_pre path. |
| 251 | template <typename scalar_t> |
| 252 | static void hc_pre_fuse_impl( |
| 253 | scalar_t* __restrict__ y, // [T, d] bf16 output |
| 254 | float* __restrict__ post_out, // [T, hc] float32 output |
| 255 | float* __restrict__ comb_out, // [T, hc*hc] float32 output |
| 256 | const scalar_t* __restrict__ x, // [T, hc, d] bf16 input |
| 257 | const float* __restrict__ hc_fn, // [mix_hc, hc_d] float32 row-major |
| 258 | const float* __restrict__ hc_scale, // [3] |
| 259 | const float* __restrict__ hc_base, // [mix_hc] |
| 260 | int64_t T, |
| 261 | int64_t d, |
| 262 | int hc, |
| 263 | int sinkhorn_iters, |
| 264 | float hc_eps, |
| 265 | float norm_eps) { |
| 266 | using bVec = at::vec::Vectorized<scalar_t>; |
| 267 | using fVec = at::vec::Vectorized<float>; |
| 268 | constexpr int64_t kVecSize = bVec::size(); // 32 bf16 |
| 269 | constexpr int64_t kFVecSize = fVec::size(); // 16 fp32 |
| 270 | constexpr int64_t kSmallTKBlockCap = 128; |
| 271 | constexpr int64_t kLargeTKBlockCap = 256; |
| 272 | constexpr int64_t kDVecsPerBlock = 16; |
| 273 | constexpr int64_t kCacheKBlockCap = kLargeTKBlockCap; |
| 274 | constexpr int64_t MAX_K_VECS = kCacheKBlockCap / kVecSize; |
| 275 | const int mix_hc = (2 + hc) * hc; |
| 276 | const int64_t hc_d = hc * d; |
| 277 | const float s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; |
| 278 | |
| 279 | const int64_t T_BLOCK = choose_t_block(T); |
| 280 | const int64_t num_threads = static_cast<int64_t>(at::get_num_threads()); |
| 281 | const int64_t K_BLOCK_CAP = (T >= num_threads) ? kLargeTKBlockCap : kSmallTKBlockCap; |
| 282 | const int64_t K_BLOCK = choose_k_block(mix_hc, kVecSize, K_BLOCK_CAP); |
| 283 | TORCH_CHECK( |
| 284 | K_BLOCK <= K_BLOCK_CAP && K_BLOCK_CAP <= kCacheKBlockCap && K_BLOCK % kVecSize == 0, |
| 285 | "hc_pre_fuse_impl: K_BLOCK must be within capacity and a multiple of vector size"); |
| 286 | const int64_t D_BLOCK = kDVecsPerBlock * kVecSize; |
| 287 | |
| 288 | const int64_t n_t_blocks = div_up(T, T_BLOCK); |
| 289 | const int64_t n_k_blocks = div_up(hc_d, K_BLOCK); |
| 290 | |
| 291 | // Partial buffer: [n_t_blocks, n_k_blocks, T_BLOCK, 1+mix_hc] |
| 292 | const int64_t partial_stride = 1 + mix_hc; |
| 293 | auto partials_tensor = at::empty({n_t_blocks * n_k_blocks * T_BLOCK * partial_stride}, at::kFloat); |
| 294 | float* const partials = partials_tensor.data_ptr<float>(); |
| 295 | |
| 296 | // Per-thread scratch: [mix_buffer[mix_hc] | pre_gate[hc] | |
| 297 | // post_gate[hc] | comb_matrix[hc*hc] | sinkhorn_scratch[3*hc]] |
| 298 | const int64_t per_thread_scratch = mix_hc + hc + hc + hc * hc + 3 * hc; |
| 299 | auto scratch_tensor = at::empty({num_threads * per_thread_scratch}, at::kFloat); |
| 300 | float* const scratch_base = scratch_tensor.data_ptr<float>(); |
| 301 | |
| 302 | // Phase 1. |
| 303 | parallel_mhc_phase1(n_t_blocks, n_k_blocks, [&](int64_t tb0, int64_t tb1, int64_t kb0, int64_t kb1) { |
| 304 | for (int64_t kb = kb0; kb < kb1; ++kb) { |
| 305 | const int64_t k0 = kb * K_BLOCK; |
| 306 | const int64_t klen = std::min(K_BLOCK, hc_d - k0); |
| 307 | |
| 308 | for (int64_t tb = tb0; tb < tb1; ++tb) { |
| 309 | const int64_t t0 = tb * T_BLOCK; |
| 310 | const int64_t tlen = std::min(T_BLOCK, T - t0); |
| 311 | float* partial_block = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride; |
| 312 | |
| 313 | const float* w_ptr[mix_hc]; |
| 314 | for (int h = 0; h < mix_hc; ++h) |
| 315 | w_ptr[h] = hc_fn + h * hc_d + k0; |
| 316 | |
| 317 | // Stack cache is sized for the largest selected K block, but only the first |
| 318 | // n_full entries are touched for the current K_BLOCK. |
| 319 | fVec x_cache_lo[MAX_K_VECS], x_cache_hi[MAX_K_VECS]; |
| 320 | |
| 321 | for (int64_t tl = 0; tl < tlen; ++tl) { |
| 322 | const scalar_t* x_t = x + (t0 + tl) * hc_d + k0; |
| 323 | float* partial_row = partial_block + tl * partial_stride; |
| 324 | |
| 325 | const int64_t n_full = klen / kVecSize; |
| 326 | const int64_t k_tail = n_full * kVecSize; |
| 327 | const int64_t rem = klen - k_tail; |
| 328 | const int64_t rem0 = std::min(rem, kFVecSize); |
| 329 | const int64_t rem1 = rem - rem0; |
| 330 | |
| 331 | // Prefetch next token. |
| 332 | if (tl + 1 < tlen) { |
| 333 | const scalar_t* xn = x + (t0 + tl + 1) * hc_d + k0; |
| 334 | __builtin_prefetch(xn, 0, 0); |
| 335 | if (klen > kVecSize * 4) __builtin_prefetch(xn + kVecSize * 4, 0, 0); |
| 336 | } |
| 337 | |
| 338 | // h=0..3: convert/cache x and accumulate sq + four dots. |
| 339 | if (tl == 0) { |
| 340 | for (int pw = 1; pw < mix_hc; ++pw) |
| 341 | __builtin_prefetch(w_ptr[pw], 0, 3); |
| 342 | } |
| 343 | { |
| 344 | const float* w0 = w_ptr[0]; |
| 345 | const float* w1 = w_ptr[1]; |
| 346 | const float* w2 = w_ptr[2]; |
| 347 | const float* w3 = w_ptr[3]; |
| 348 | fVec sq_lo(0.f), sq_hi(0.f); |
| 349 | fVec d0_lo(0.f), d0_hi(0.f); |
| 350 | fVec d1_lo(0.f), d1_hi(0.f); |
| 351 | fVec d2_lo(0.f), d2_hi(0.f); |
| 352 | fVec d3_lo(0.f), d3_hi(0.f); |
| 353 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 354 | const int64_t k = ki * kVecSize; |
| 355 | std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k)); |
| 356 | sq_lo += x_cache_lo[ki] * x_cache_lo[ki]; |
| 357 | sq_hi += x_cache_hi[ki] * x_cache_hi[ki]; |
| 358 | d0_lo += x_cache_lo[ki] * fVec::loadu(w0 + k); |
| 359 | d0_hi += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize); |
| 360 | d1_lo += x_cache_lo[ki] * fVec::loadu(w1 + k); |
| 361 | d1_hi += x_cache_hi[ki] * fVec::loadu(w1 + k + kFVecSize); |
| 362 | d2_lo += x_cache_lo[ki] * fVec::loadu(w2 + k); |
| 363 | d2_hi += x_cache_hi[ki] * fVec::loadu(w2 + k + kFVecSize); |
| 364 | d3_lo += x_cache_lo[ki] * fVec::loadu(w3 + k); |
| 365 | d3_hi += x_cache_hi[ki] * fVec::loadu(w3 + k + kFVecSize); |
| 366 | } |
| 367 | |
| 368 | fVec x_tail_lo(0.f), x_tail_hi(0.f); |
| 369 | if (rem > 0) { |
| 370 | std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem)); |
| 371 | sq_lo += x_tail_lo * x_tail_lo; |
| 372 | if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi; |
| 373 | d0_lo += x_tail_lo * fVec::loadu(w0 + k_tail, rem0); |
| 374 | d1_lo += x_tail_lo * fVec::loadu(w1 + k_tail, rem0); |
| 375 | d2_lo += x_tail_lo * fVec::loadu(w2 + k_tail, rem0); |
| 376 | d3_lo += x_tail_lo * fVec::loadu(w3 + k_tail, rem0); |
| 377 | if (rem1 > 0) { |
| 378 | d0_hi += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1); |
| 379 | d1_hi += x_tail_hi * fVec::loadu(w1 + k_tail + kFVecSize, rem1); |
| 380 | d2_hi += x_tail_hi * fVec::loadu(w2 + k_tail + kFVecSize, rem1); |
| 381 | d3_hi += x_tail_hi * fVec::loadu(w3 + k_tail + kFVecSize, rem1); |
| 382 | } |
| 383 | } |
| 384 | |
| 385 | partial_row[0] = vec_reduce_sum(sq_lo + sq_hi); |
| 386 | partial_row[1 + 0] = vec_reduce_sum(d0_lo + d0_hi); |
| 387 | partial_row[1 + 1] = vec_reduce_sum(d1_lo + d1_hi); |
| 388 | partial_row[1 + 2] = vec_reduce_sum(d2_lo + d2_hi); |
| 389 | partial_row[1 + 3] = vec_reduce_sum(d3_lo + d3_hi); |
| 390 | |
| 391 | // h=4..mix_hc-1: reuse cached x, four dots per pass. |
| 392 | int h = 4; |
| 393 | for (; h + 3 < mix_hc; h += 4) { |
| 394 | const float* wa = w_ptr[h]; |
| 395 | const float* wb = w_ptr[h + 1]; |
| 396 | const float* wc = w_ptr[h + 2]; |
| 397 | const float* wd = w_ptr[h + 3]; |
| 398 | fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); |
| 399 | fVec c0(0.f), c1(0.f), d0(0.f), d1(0.f); |
| 400 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 401 | const int64_t k = ki * kVecSize; |
| 402 | a0 += x_cache_lo[ki] * fVec::loadu(wa + k); |
| 403 | b0 += x_cache_lo[ki] * fVec::loadu(wb + k); |
| 404 | c0 += x_cache_lo[ki] * fVec::loadu(wc + k); |
| 405 | d0 += x_cache_lo[ki] * fVec::loadu(wd + k); |
| 406 | a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); |
| 407 | b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); |
| 408 | c1 += x_cache_hi[ki] * fVec::loadu(wc + k + kFVecSize); |
| 409 | d1 += x_cache_hi[ki] * fVec::loadu(wd + k + kFVecSize); |
| 410 | } |
| 411 | if (rem > 0) { |
| 412 | a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); |
| 413 | b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); |
| 414 | c0 += x_tail_lo * fVec::loadu(wc + k_tail, rem0); |
| 415 | d0 += x_tail_lo * fVec::loadu(wd + k_tail, rem0); |
| 416 | if (rem1 > 0) { |
| 417 | a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); |
| 418 | b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); |
| 419 | c1 += x_tail_hi * fVec::loadu(wc + k_tail + kFVecSize, rem1); |
| 420 | d1 += x_tail_hi * fVec::loadu(wd + k_tail + kFVecSize, rem1); |
| 421 | } |
| 422 | } |
| 423 | partial_row[1 + h] = vec_reduce_sum(a0 + a1); |
| 424 | partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); |
| 425 | partial_row[1 + h + 2] = vec_reduce_sum(c0 + c1); |
| 426 | partial_row[1 + h + 3] = vec_reduce_sum(d0 + d1); |
| 427 | } |
| 428 | for (; h + 1 < mix_hc; h += 2) { |
| 429 | const float* wa = w_ptr[h]; |
| 430 | const float* wb = w_ptr[h + 1]; |
| 431 | fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); |
| 432 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 433 | const int64_t k = ki * kVecSize; |
| 434 | a0 += x_cache_lo[ki] * fVec::loadu(wa + k); |
| 435 | b0 += x_cache_lo[ki] * fVec::loadu(wb + k); |
| 436 | a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); |
| 437 | b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); |
| 438 | } |
| 439 | if (rem > 0) { |
| 440 | a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); |
| 441 | b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); |
| 442 | if (rem1 > 0) { |
| 443 | a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); |
| 444 | b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); |
| 445 | } |
| 446 | } |
| 447 | partial_row[1 + h] = vec_reduce_sum(a0 + a1); |
| 448 | partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); |
| 449 | } |
| 450 | // odd leftover |
| 451 | if (h < mix_hc) { |
| 452 | const float* w = w_ptr[h]; |
| 453 | fVec hd0(0.f), hd1(0.f); |
| 454 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 455 | const int64_t k = ki * kVecSize; |
| 456 | hd0 += x_cache_lo[ki] * fVec::loadu(w + k); |
| 457 | hd1 += x_cache_hi[ki] * fVec::loadu(w + k + kFVecSize); |
| 458 | } |
| 459 | if (rem > 0) { |
| 460 | hd0 += x_tail_lo * fVec::loadu(w + k_tail, rem0); |
| 461 | if (rem1 > 0) hd1 += x_tail_hi * fVec::loadu(w + k_tail + kFVecSize, rem1); |
| 462 | } |
| 463 | partial_row[1 + h] = vec_reduce_sum(hd0 + hd1); |
| 464 | } |
| 465 | } |
| 466 | } // tl |
| 467 | |
| 468 | // Prefetch next t-block head. |
| 469 | if (tb + 1 < tb1) { |
| 470 | const scalar_t* x_next_tb = x + ((tb + 1) * T_BLOCK) * hc_d + k0; |
| 471 | __builtin_prefetch(x_next_tb, 0, 0); |
| 472 | if (klen > kVecSize * 4) __builtin_prefetch(x_next_tb + kVecSize * 4, 0, 0); |
| 473 | } |
| 474 | } // tb |
| 475 | } // kb |
| 476 | }); // Phase 1 |
| 477 | |
| 478 | // Phase 2+3: reduce partials, sinkhorn, and combine. |
| 479 | at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) { |
| 480 | const int64_t tid = at::get_thread_num(); |
| 481 | fVec v0, v1, pre_fvec; |
| 482 | float* const thread_scratch = scratch_base + tid * per_thread_scratch; |
| 483 | float* const mix_buffer = thread_scratch; |
| 484 | float* const pre_gate = thread_scratch + mix_hc; |
| 485 | float* const post_gate = pre_gate + hc; |
| 486 | float* const comb_matrix = post_gate + hc; |
| 487 | float* const sinkhorn_scratch = comb_matrix + hc * hc; |
| 488 | |
| 489 | for (int64_t t = begin; t < end; ++t) { |
| 490 | const int64_t tb = t / T_BLOCK; |
| 491 | const int64_t tl = t % T_BLOCK; |
| 492 | |
| 493 | std::fill(mix_buffer, mix_buffer + mix_hc, 0.f); |
| 494 | double sq_total = 0.0; |
| 495 | for (int64_t kb = 0; kb < n_k_blocks; ++kb) { |
| 496 | const float* partial_row = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride + tl * partial_stride; |
| 497 | sq_total += static_cast<double>(partial_row[0]); |
| 498 | #pragma GCC unroll 4 |
| 499 | for (int i = 0; i < mix_hc; ++i) |
| 500 | mix_buffer[i] += partial_row[1 + i]; |
| 501 | } |
| 502 | const float inv_rms = |
| 503 | static_cast<float>(1.0 / std::sqrt(sq_total / static_cast<double>(hc_d) + static_cast<double>(norm_eps))); |
| 504 | #pragma GCC unroll 4 |
| 505 | for (int i = 0; i < mix_hc; ++i) |
| 506 | mix_buffer[i] *= inv_rms; |
| 507 | |
| 508 | parse_mixes_and_sinkhorn( |
| 509 | pre_gate, |
| 510 | post_gate, |
| 511 | comb_matrix, |
| 512 | mix_buffer, |
| 513 | s0, |
| 514 | s1, |
| 515 | s2, |
| 516 | hc_base, |
| 517 | sinkhorn_iters, |
| 518 | hc_eps, |
| 519 | hc, |
| 520 | sinkhorn_scratch); |
| 521 | |
| 522 | std::copy_n(post_gate, hc, post_out + t * hc); |
| 523 | std::copy_n(comb_matrix, hc * hc, comb_out + t * hc * hc); |
| 524 | |
| 525 | // Phase 3: d-tiled combine. |
| 526 | const scalar_t* x_t = x + t * hc_d; |
| 527 | scalar_t* y_t = y + t * d; |
| 528 | for (int64_t j0 = 0; j0 < d; j0 += D_BLOCK) { |
| 529 | const int64_t jlen = std::min(D_BLOCK, d - j0); |
| 530 | |
| 531 | int64_t kk = 0; |
| 532 | for (; kk <= jlen - kVecSize; kk += kVecSize) { |
| 533 | pre_fvec = fVec(pre_gate[0]); |
| 534 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk)); |
| 535 | fVec acc0 = pre_fvec * v0; |
| 536 | fVec acc1 = pre_fvec * v1; |
| 537 | #pragma GCC unroll 4 |
| 538 | for (int h = 1; h < hc; ++h) { |
| 539 | pre_fvec = fVec(pre_gate[h]); |
| 540 | const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; |
| 541 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th)); |
| 542 | acc0 += pre_fvec * v0; |
| 543 | acc1 += pre_fvec * v1; |
| 544 | } |
| 545 | at::vec::convert_from_float<scalar_t>(acc0, acc1).store(y_t + j0 + kk); |
| 546 | } |
| 547 | if (kk < jlen) { |
| 548 | const int64_t rem = jlen - kk, rem1 = rem - std::min(rem, kFVecSize); |
| 549 | pre_fvec = fVec(pre_gate[0]); |
| 550 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk, rem)); |
| 551 | fVec acc0 = pre_fvec * v0; |
| 552 | fVec acc1 = pre_fvec * v1; |
| 553 | #pragma GCC unroll 4 |
| 554 | for (int h = 1; h < hc; ++h) { |
| 555 | pre_fvec = fVec(pre_gate[h]); |
| 556 | const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; |
| 557 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th, rem)); |
| 558 | acc0 += pre_fvec * v0; |
| 559 | acc1 += pre_fvec * v1; |
| 560 | } |
| 561 | at::vec::convert_from_float<scalar_t>(acc0, rem1 > 0 ? acc1 : fVec(0.f)).store(y_t + j0 + kk, rem); |
| 562 | } |
| 563 | } // d-tile |
| 564 | } // t |
| 565 | }); // Phase 2+3 |
| 566 | } |
| 567 | |
| 568 | // Fused hc_head path. |
| 569 | template <typename scalar_t> |
| 570 | static void hc_head_fuse_impl( |
| 571 | scalar_t* __restrict__ y, |
| 572 | const scalar_t* __restrict__ x, |
| 573 | const float* __restrict__ hc_fn, // [hc, hc*d] row-major |
| 574 | float hc_scale_val, |
| 575 | const float* __restrict__ hc_base, |
| 576 | int64_t T, |
| 577 | int64_t d, |
| 578 | int hc, |
| 579 | float hc_eps, |
| 580 | float norm_eps) { |
| 581 | using bVec = at::vec::Vectorized<scalar_t>; |
| 582 | using fVec = at::vec::Vectorized<float>; |
| 583 | constexpr int64_t kVecSize = bVec::size(); // 32 bf16 |
| 584 | constexpr int64_t kFVecSize = fVec::size(); // 16 fp32 |
| 585 | constexpr int64_t kDVecsPerBlock = 16; |
| 586 | const int64_t hc_d = hc * d; |
| 587 | |
| 588 | constexpr int64_t K_BLOCK_CAP = 1024; |
| 589 | constexpr int64_t MAX_K_VECS = K_BLOCK_CAP / kVecSize; |
| 590 | |
| 591 | const int64_t T_BLOCK = choose_t_block(T); |
| 592 | const int64_t num_threads = static_cast<int64_t>(at::get_num_threads()); |
| 593 | const int64_t K_BLOCK = choose_k_block(hc, kVecSize, K_BLOCK_CAP); |
| 594 | TORCH_CHECK( |
| 595 | K_BLOCK <= K_BLOCK_CAP && K_BLOCK % kVecSize == 0, |
| 596 | "hc_head_fuse_impl: K_BLOCK must be <= K_BLOCK_CAP and a multiple of vector size"); |
| 597 | const int64_t D_BLOCK = kDVecsPerBlock * kVecSize; |
| 598 | |
| 599 | const int64_t n_t_blocks = div_up(T, T_BLOCK); |
| 600 | const int64_t n_k_blocks = div_up(hc_d, K_BLOCK); |
| 601 | |
| 602 | // Partial buffer: [n_t_blocks, n_k_blocks, T_BLOCK, 1+hc] |
| 603 | // p[0] = partial sq, p[1..hc] = partial dots |
| 604 | const int64_t partial_stride = 1 + hc; |
| 605 | auto partials_tensor = at::empty({n_t_blocks * n_k_blocks * T_BLOCK * partial_stride}, at::kFloat); |
| 606 | float* const partials = partials_tensor.data_ptr<float>(); |
| 607 | |
| 608 | // Per-thread scratch: [pre_gate[hc]] |
| 609 | const int64_t per_thread_head = hc; |
| 610 | auto scratch_tensor = at::empty({num_threads * per_thread_head}, at::kFloat); |
| 611 | float* const scratch_base = scratch_tensor.data_ptr<float>(); |
| 612 | |
| 613 | // Phase 1. |
| 614 | parallel_mhc_phase1(n_t_blocks, n_k_blocks, [&](int64_t tb0, int64_t tb1, int64_t kb0, int64_t kb1) { |
| 615 | for (int64_t kb = kb0; kb < kb1; ++kb) { |
| 616 | const int64_t k0 = kb * K_BLOCK; |
| 617 | const int64_t klen = std::min(K_BLOCK, hc_d - k0); |
| 618 | |
| 619 | for (int64_t tb = tb0; tb < tb1; ++tb) { |
| 620 | const int64_t t0 = tb * T_BLOCK; |
| 621 | const int64_t tlen = std::min(T_BLOCK, T - t0); |
| 622 | float* partial_block = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride; |
| 623 | |
| 624 | const float* w_ptr[hc]; |
| 625 | for (int h = 0; h < hc; ++h) |
| 626 | w_ptr[h] = hc_fn + h * hc_d + k0; |
| 627 | |
| 628 | // Stack cache is sized for the largest selected K block, but only the first |
| 629 | // n_full entries are touched for the current K_BLOCK. |
| 630 | fVec x_cache_lo[MAX_K_VECS], x_cache_hi[MAX_K_VECS]; |
| 631 | |
| 632 | for (int64_t tl = 0; tl < tlen; ++tl) { |
| 633 | const scalar_t* x_t = x + (t0 + tl) * hc_d + k0; |
| 634 | float* partial_row = partial_block + tl * partial_stride; |
| 635 | |
| 636 | const int64_t n_full = klen / kVecSize; |
| 637 | const int64_t k_tail = n_full * kVecSize; |
| 638 | const int64_t rem = klen - k_tail; |
| 639 | const int64_t rem0 = std::min(rem, kFVecSize); |
| 640 | const int64_t rem1 = rem - rem0; |
| 641 | |
| 642 | // Prefetch next token. |
| 643 | if (tl + 1 < tlen) { |
| 644 | const scalar_t* xn = x + (t0 + tl + 1) * hc_d + k0; |
| 645 | __builtin_prefetch(xn, 0, 0); |
| 646 | if (klen > kVecSize * 4) __builtin_prefetch(xn + kVecSize * 4, 0, 0); |
| 647 | } |
| 648 | |
| 649 | // Convert/cache x and accumulate sq + multiple dots. |
| 650 | if (tl == 0) { |
| 651 | for (int pw = 1; pw < hc; ++pw) |
| 652 | __builtin_prefetch(w_ptr[pw], 0, 3); |
| 653 | } |
| 654 | { |
| 655 | fVec x_tail_lo(0.f), x_tail_hi(0.f); |
| 656 | int h = 0; |
| 657 | |
| 658 | if (hc >= 4) { |
| 659 | const float* w0 = w_ptr[0]; |
| 660 | const float* w1 = w_ptr[1]; |
| 661 | const float* w2 = w_ptr[2]; |
| 662 | const float* w3 = w_ptr[3]; |
| 663 | fVec sq_lo(0.f), sq_hi(0.f); |
| 664 | fVec d0_lo(0.f), d0_hi(0.f); |
| 665 | fVec d1_lo(0.f), d1_hi(0.f); |
| 666 | fVec d2_lo(0.f), d2_hi(0.f); |
| 667 | fVec d3_lo(0.f), d3_hi(0.f); |
| 668 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 669 | const int64_t k = ki * kVecSize; |
| 670 | std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k)); |
| 671 | sq_lo += x_cache_lo[ki] * x_cache_lo[ki]; |
| 672 | sq_hi += x_cache_hi[ki] * x_cache_hi[ki]; |
| 673 | d0_lo += x_cache_lo[ki] * fVec::loadu(w0 + k); |
| 674 | d0_hi += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize); |
| 675 | d1_lo += x_cache_lo[ki] * fVec::loadu(w1 + k); |
| 676 | d1_hi += x_cache_hi[ki] * fVec::loadu(w1 + k + kFVecSize); |
| 677 | d2_lo += x_cache_lo[ki] * fVec::loadu(w2 + k); |
| 678 | d2_hi += x_cache_hi[ki] * fVec::loadu(w2 + k + kFVecSize); |
| 679 | d3_lo += x_cache_lo[ki] * fVec::loadu(w3 + k); |
| 680 | d3_hi += x_cache_hi[ki] * fVec::loadu(w3 + k + kFVecSize); |
| 681 | } |
| 682 | if (rem > 0) { |
| 683 | std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem)); |
| 684 | sq_lo += x_tail_lo * x_tail_lo; |
| 685 | if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi; |
| 686 | d0_lo += x_tail_lo * fVec::loadu(w0 + k_tail, rem0); |
| 687 | d1_lo += x_tail_lo * fVec::loadu(w1 + k_tail, rem0); |
| 688 | d2_lo += x_tail_lo * fVec::loadu(w2 + k_tail, rem0); |
| 689 | d3_lo += x_tail_lo * fVec::loadu(w3 + k_tail, rem0); |
| 690 | if (rem1 > 0) { |
| 691 | d0_hi += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1); |
| 692 | d1_hi += x_tail_hi * fVec::loadu(w1 + k_tail + kFVecSize, rem1); |
| 693 | d2_hi += x_tail_hi * fVec::loadu(w2 + k_tail + kFVecSize, rem1); |
| 694 | d3_hi += x_tail_hi * fVec::loadu(w3 + k_tail + kFVecSize, rem1); |
| 695 | } |
| 696 | } |
| 697 | partial_row[0] = vec_reduce_sum(sq_lo + sq_hi); |
| 698 | partial_row[1 + 0] = vec_reduce_sum(d0_lo + d0_hi); |
| 699 | partial_row[1 + 1] = vec_reduce_sum(d1_lo + d1_hi); |
| 700 | partial_row[1 + 2] = vec_reduce_sum(d2_lo + d2_hi); |
| 701 | partial_row[1 + 3] = vec_reduce_sum(d3_lo + d3_hi); |
| 702 | h = 4; |
| 703 | } else { |
| 704 | const float* w0 = w_ptr[0]; |
| 705 | fVec sq_lo(0.f), sq_hi(0.f); |
| 706 | fVec d0(0.f), d1(0.f); |
| 707 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 708 | const int64_t k = ki * kVecSize; |
| 709 | std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k)); |
| 710 | sq_lo += x_cache_lo[ki] * x_cache_lo[ki]; |
| 711 | sq_hi += x_cache_hi[ki] * x_cache_hi[ki]; |
| 712 | d0 += x_cache_lo[ki] * fVec::loadu(w0 + k); |
| 713 | d1 += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize); |
| 714 | } |
| 715 | if (rem > 0) { |
| 716 | std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem)); |
| 717 | sq_lo += x_tail_lo * x_tail_lo; |
| 718 | if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi; |
| 719 | d0 += x_tail_lo * fVec::loadu(w0 + k_tail, rem0); |
| 720 | if (rem1 > 0) d1 += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1); |
| 721 | } |
| 722 | partial_row[0] = vec_reduce_sum(sq_lo + sq_hi); |
| 723 | partial_row[1 + 0] = vec_reduce_sum(d0 + d1); |
| 724 | h = 1; |
| 725 | } |
| 726 | |
| 727 | // h: reuse cached x, four dots per pass. |
| 728 | for (; h + 3 < hc; h += 4) { |
| 729 | const float* wa = w_ptr[h]; |
| 730 | const float* wb = w_ptr[h + 1]; |
| 731 | const float* wc = w_ptr[h + 2]; |
| 732 | const float* wd = w_ptr[h + 3]; |
| 733 | fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); |
| 734 | fVec c0(0.f), c1(0.f), d0(0.f), d1(0.f); |
| 735 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 736 | const int64_t k = ki * kVecSize; |
| 737 | a0 += x_cache_lo[ki] * fVec::loadu(wa + k); |
| 738 | b0 += x_cache_lo[ki] * fVec::loadu(wb + k); |
| 739 | c0 += x_cache_lo[ki] * fVec::loadu(wc + k); |
| 740 | d0 += x_cache_lo[ki] * fVec::loadu(wd + k); |
| 741 | a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); |
| 742 | b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); |
| 743 | c1 += x_cache_hi[ki] * fVec::loadu(wc + k + kFVecSize); |
| 744 | d1 += x_cache_hi[ki] * fVec::loadu(wd + k + kFVecSize); |
| 745 | } |
| 746 | if (rem > 0) { |
| 747 | a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); |
| 748 | b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); |
| 749 | c0 += x_tail_lo * fVec::loadu(wc + k_tail, rem0); |
| 750 | d0 += x_tail_lo * fVec::loadu(wd + k_tail, rem0); |
| 751 | if (rem1 > 0) { |
| 752 | a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); |
| 753 | b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); |
| 754 | c1 += x_tail_hi * fVec::loadu(wc + k_tail + kFVecSize, rem1); |
| 755 | d1 += x_tail_hi * fVec::loadu(wd + k_tail + kFVecSize, rem1); |
| 756 | } |
| 757 | } |
| 758 | partial_row[1 + h] = vec_reduce_sum(a0 + a1); |
| 759 | partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); |
| 760 | partial_row[1 + h + 2] = vec_reduce_sum(c0 + c1); |
| 761 | partial_row[1 + h + 3] = vec_reduce_sum(d0 + d1); |
| 762 | } |
| 763 | for (; h + 1 < hc; h += 2) { |
| 764 | const float* wa = w_ptr[h]; |
| 765 | const float* wb = w_ptr[h + 1]; |
| 766 | fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); |
| 767 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 768 | const int64_t k = ki * kVecSize; |
| 769 | a0 += x_cache_lo[ki] * fVec::loadu(wa + k); |
| 770 | b0 += x_cache_lo[ki] * fVec::loadu(wb + k); |
| 771 | a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); |
| 772 | b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); |
| 773 | } |
| 774 | if (rem > 0) { |
| 775 | a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); |
| 776 | b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); |
| 777 | if (rem1 > 0) { |
| 778 | a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); |
| 779 | b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); |
| 780 | } |
| 781 | } |
| 782 | partial_row[1 + h] = vec_reduce_sum(a0 + a1); |
| 783 | partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); |
| 784 | } |
| 785 | // odd leftover |
| 786 | if (h < hc) { |
| 787 | const float* w = w_ptr[h]; |
| 788 | fVec hd0(0.f), hd1(0.f); |
| 789 | for (int64_t ki = 0; ki < n_full; ++ki) { |
| 790 | const int64_t k = ki * kVecSize; |
| 791 | hd0 += x_cache_lo[ki] * fVec::loadu(w + k); |
| 792 | hd1 += x_cache_hi[ki] * fVec::loadu(w + k + kFVecSize); |
| 793 | } |
| 794 | if (rem > 0) { |
| 795 | hd0 += x_tail_lo * fVec::loadu(w + k_tail, rem0); |
| 796 | if (rem1 > 0) hd1 += x_tail_hi * fVec::loadu(w + k_tail + kFVecSize, rem1); |
| 797 | } |
| 798 | partial_row[1 + h] = vec_reduce_sum(hd0 + hd1); |
| 799 | } |
| 800 | } |
| 801 | } // tl |
| 802 | |
| 803 | // Prefetch next t-block head. |
| 804 | if (tb + 1 < tb1) { |
| 805 | const scalar_t* x_next_tb = x + ((tb + 1) * T_BLOCK) * hc_d + k0; |
| 806 | __builtin_prefetch(x_next_tb, 0, 0); |
| 807 | if (klen > kVecSize * 2) __builtin_prefetch(x_next_tb + kVecSize * 2, 0, 0); |
| 808 | } |
| 809 | } // tb |
| 810 | } // kb |
| 811 | }); // Phase 1 |
| 812 | |
| 813 | // Phase 2+3. |
| 814 | at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) { |
| 815 | const int64_t tid = at::get_thread_num(); |
| 816 | float* const pre_gate = scratch_base + tid * per_thread_head; |
| 817 | fVec v0, v1, pre_fvec; |
| 818 | |
| 819 | for (int64_t t = begin; t < end; ++t) { |
| 820 | const int64_t tb = t / T_BLOCK; |
| 821 | const int64_t tl = t % T_BLOCK; |
| 822 | |
| 823 | // Phase 2: reduce partials + sigmoid gate. |
| 824 | double sq_total = 0.0; |
| 825 | std::fill(pre_gate, pre_gate + hc, 0.f); |
| 826 | for (int64_t kb = 0; kb < n_k_blocks; ++kb) { |
| 827 | const float* partial_row = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride + tl * partial_stride; |
| 828 | sq_total += static_cast<double>(partial_row[0]); |
| 829 | #pragma GCC unroll 4 |
| 830 | for (int h = 0; h < hc; ++h) |
| 831 | pre_gate[h] += partial_row[1 + h]; |
| 832 | } |
| 833 | const float inv_rms = |
| 834 | static_cast<float>(1.0 / std::sqrt(sq_total / static_cast<double>(hc_d) + static_cast<double>(norm_eps))); |
| 835 | #pragma GCC unroll 4 |
| 836 | for (int h = 0; h < hc; ++h) { |
| 837 | const float gate = pre_gate[h] * inv_rms * hc_scale_val + hc_base[h]; |
| 838 | pre_gate[h] = 1.f / (1.f + std::exp(-gate)) + hc_eps; |
| 839 | } |
| 840 | |
| 841 | // Phase 3: d-tiled combine. |
| 842 | const scalar_t* x_t = x + t * hc_d; |
| 843 | scalar_t* y_t = y + t * d; |
| 844 | for (int64_t j0 = 0; j0 < d; j0 += D_BLOCK) { |
| 845 | const int64_t jlen = std::min(D_BLOCK, d - j0); |
| 846 | |
| 847 | int64_t kk = 0; |
| 848 | for (; kk <= jlen - kVecSize; kk += kVecSize) { |
| 849 | pre_fvec = fVec(pre_gate[0]); |
| 850 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk)); |
| 851 | fVec acc0 = pre_fvec * v0; |
| 852 | fVec acc1 = pre_fvec * v1; |
| 853 | #pragma GCC unroll 4 |
| 854 | for (int h = 1; h < hc; ++h) { |
| 855 | pre_fvec = fVec(pre_gate[h]); |
| 856 | const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; |
| 857 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th)); |
| 858 | acc0 += pre_fvec * v0; |
| 859 | acc1 += pre_fvec * v1; |
| 860 | } |
| 861 | at::vec::convert_from_float<scalar_t>(acc0, acc1).store(y_t + j0 + kk); |
| 862 | } |
| 863 | if (kk < jlen) { |
| 864 | const int64_t rem = jlen - kk, rem1 = rem - std::min(rem, kFVecSize); |
| 865 | pre_fvec = fVec(pre_gate[0]); |
| 866 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk, rem)); |
| 867 | fVec acc0 = pre_fvec * v0; |
| 868 | fVec acc1 = pre_fvec * v1; |
| 869 | #pragma GCC unroll 4 |
| 870 | for (int h = 1; h < hc; ++h) { |
| 871 | pre_fvec = fVec(pre_gate[h]); |
| 872 | const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; |
| 873 | std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th, rem)); |
| 874 | acc0 += pre_fvec * v0; |
| 875 | acc1 += pre_fvec * v1; |
| 876 | } |
| 877 | at::vec::convert_from_float<scalar_t>(acc0, rem1 > 0 ? acc1 : fVec(0.f)).store(y_t + j0 + kk, rem); |
| 878 | } |
| 879 | } // d-tile |
| 880 | } // t |
| 881 | }); // Phase 2+3 |
| 882 | } |
| 883 | |
| 884 | } // anonymous namespace |
| 885 | |
| 886 | // Fused hc_pre entry. |
| 887 | std::tuple<at::Tensor, at::Tensor, at::Tensor> hc_pre_fused_cpu( |
| 888 | at::Tensor& x, |
| 889 | at::Tensor& hc_fn, |
| 890 | at::Tensor& hc_scale, |
| 891 | at::Tensor& hc_base, |
| 892 | int64_t hc_mult, |
| 893 | int64_t sinkhorn_iters, |
| 894 | double rms_eps, |
| 895 | double hc_eps) { |
| 896 | TORCH_CHECK(hc_mult >= 2, "hc_pre_fused_cpu: hc_mult must be >= 2"); |
| 897 | TORCH_CHECK( |
| 898 | x.dim() == 3 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf), |
| 899 | "hc_pre_fused_cpu: x must be bf16/fp16 [T, hc, d]"); |
| 900 | TORCH_CHECK(hc_fn.scalar_type() == at::kFloat, "hc_pre_fused_cpu: hc_fn must be float32"); |
| 901 | TORCH_CHECK( |
| 902 | hc_scale.scalar_type() == at::kFloat && hc_scale.numel() == 3, "hc_pre_fused_cpu: hc_scale must be float32 [3]"); |
| 903 | TORCH_CHECK(hc_base.scalar_type() == at::kFloat, "hc_pre_fused_cpu: hc_base must be float32"); |
| 904 | TORCH_CHECK( |
| 905 | x.device().is_cpu() && hc_fn.device().is_cpu() && hc_scale.device().is_cpu() && hc_base.device().is_cpu(), |
| 906 | "hc_pre_fused_cpu: all inputs must be CPU tensors"); |
| 907 | |
| 908 | const int64_t T = x.size(0); |
| 909 | const int64_t d = x.size(2); |
| 910 | const int64_t hc_d = hc_mult * d; |
| 911 | const int64_t mix_hc = (2 + hc_mult) * hc_mult; |
| 912 | |
| 913 | TORCH_CHECK(T > 0 && d > 0, "hc_pre_fused_cpu: T and d must be positive"); |
| 914 | TORCH_CHECK(x.size(1) == hc_mult, "hc_pre_fused_cpu: x shape must be [T, hc_mult, d]"); |
| 915 | TORCH_CHECK( |
| 916 | hc_fn.dim() == 2 && hc_fn.size(0) == mix_hc && hc_fn.size(1) == hc_d, |
| 917 | "hc_pre_fused_cpu: hc_fn must be [mix_hc, hc_mult*d]"); |
| 918 | TORCH_CHECK(hc_base.numel() == mix_hc, "hc_pre_fused_cpu: hc_base.numel() must equal mix_hc"); |
| 919 | |
| 920 | // Static per-layer weights (hc_fn/hc_scale/hc_base) are allocated fp32 and |
| 921 | // never reassigned after weight loading; x/residual come straight from a |
| 922 | // preceding op. All are expected contiguous already -- fail loudly instead |
| 923 | // of silently absorbing a copy if that contract is ever violated. |
| 924 | TORCH_CHECK( |
| 925 | x.is_contiguous() && hc_fn.is_contiguous() && hc_scale.is_contiguous() && hc_base.is_contiguous(), |
| 926 | "hc_pre_fused_cpu: x/hc_fn/hc_scale/hc_base must be contiguous"); |
| 927 | |
| 928 | auto f32_opts = at::TensorOptions().dtype(at::kFloat).device(x.device()); |
| 929 | auto post = at::empty({T, hc_mult}, f32_opts); |
| 930 | auto comb = at::empty({T, hc_mult, hc_mult}, f32_opts); |
| 931 | auto y = at::empty({T, d}, x.options()); |
| 932 | |
| 933 | AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_pre_fused_cpu", [&] { |
| 934 | hc_pre_fuse_impl<scalar_t>( |
| 935 | y.data_ptr<scalar_t>(), |
| 936 | post.data_ptr<float>(), |
| 937 | comb.data_ptr<float>(), |
| 938 | x.data_ptr<scalar_t>(), |
| 939 | hc_fn.data_ptr<float>(), |
| 940 | hc_scale.data_ptr<float>(), |
| 941 | hc_base.data_ptr<float>(), |
| 942 | T, |
| 943 | d, |
| 944 | static_cast<int>(hc_mult), |
| 945 | static_cast<int>(sinkhorn_iters), |
| 946 | static_cast<float>(hc_eps), |
| 947 | static_cast<float>(rms_eps)); |
| 948 | }); |
| 949 | return {y, post, comb}; |
| 950 | } |
| 951 | |
| 952 | // Fused hc_post entry. |
| 953 | at::Tensor hc_post_fused_cpu(at::Tensor& x, at::Tensor& residual, at::Tensor& post, at::Tensor& comb) { |
| 954 | TORCH_CHECK( |
| 955 | x.dim() == 2 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf), |
| 956 | "hc_post_fused_cpu: x must be bf16/fp16 [T, d]"); |
| 957 | TORCH_CHECK( |
| 958 | residual.dim() == 3 && residual.scalar_type() == x.scalar_type(), |
| 959 | "hc_post_fused_cpu: residual must be same dtype as x [T, hc, d]"); |
| 960 | TORCH_CHECK(post.scalar_type() == at::kFloat, "hc_post_fused_cpu: post must be float32"); |
| 961 | TORCH_CHECK(comb.scalar_type() == at::kFloat, "hc_post_fused_cpu: comb must be float32"); |
| 962 | TORCH_CHECK( |
| 963 | x.device().is_cpu() && residual.device().is_cpu() && post.device().is_cpu() && comb.device().is_cpu(), |
| 964 | "hc_post_fused_cpu: all inputs must be CPU tensors"); |
| 965 | |
| 966 | const int64_t T = x.size(0); |
| 967 | const int64_t d = x.size(1); |
| 968 | const int64_t hc = residual.size(1); |
| 969 | TORCH_CHECK(hc >= 2, "hc_post_fused_cpu: hc_mult must be >= 2"); |
| 970 | TORCH_CHECK(T > 0 && d > 0, "hc_post_fused_cpu: T and d must be positive"); |
| 971 | TORCH_CHECK(residual.size(0) == T && residual.size(2) == d, "hc_post_fused_cpu: residual shape must be [T, hc, d]"); |
| 972 | TORCH_CHECK( |
| 973 | post.dim() == 2 && post.size(0) == T && post.size(1) == hc, "hc_post_fused_cpu: post shape must be [T, hc]"); |
| 974 | TORCH_CHECK( |
| 975 | comb.dim() == 3 && comb.size(0) == T && comb.size(1) == hc && comb.size(2) == hc, |
| 976 | "hc_post_fused_cpu: comb shape must be [T, hc, hc]"); |
| 977 | |
| 978 | TORCH_CHECK( |
| 979 | x.is_contiguous() && residual.is_contiguous() && post.is_contiguous() && comb.is_contiguous(), |
| 980 | "hc_post_fused_cpu: x/residual/post/comb must be contiguous"); |
| 981 | |
| 982 | auto out = at::empty({T, hc, d}, x.options()); // same dtype as x |
| 983 | |
| 984 | AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_post_fused_cpu", [&] { |
| 985 | hc_post_fuse_impl<scalar_t>( |
| 986 | out.data_ptr<scalar_t>(), |
| 987 | x.data_ptr<scalar_t>(), |
| 988 | residual.data_ptr<scalar_t>(), |
| 989 | post.data_ptr<float>(), |
| 990 | comb.data_ptr<float>(), |
| 991 | T, |
| 992 | d, |
| 993 | static_cast<int>(hc)); |
| 994 | }); |
| 995 | |
| 996 | return out; |
| 997 | } |
| 998 | |
| 999 | // Fused hc_head entry. |
| 1000 | at::Tensor hc_head_fused_cpu( |
| 1001 | at::Tensor& x, at::Tensor& hc_fn, at::Tensor& hc_scale, at::Tensor& hc_base, double hc_eps, double norm_eps) { |
| 1002 | TORCH_CHECK( |
| 1003 | x.dim() == 3 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf), |
| 1004 | "hc_head_fused_cpu: x must be bf16/fp16 [T, hc, d]"); |
| 1005 | const int64_t hc_mult = x.size(1); |
| 1006 | TORCH_CHECK(hc_mult >= 2, "hc_head_fused_cpu: hc_mult must be >= 2"); |
| 1007 | TORCH_CHECK(hc_fn.scalar_type() == at::kFloat, "hc_head_fused_cpu: hc_fn must be float32"); |
| 1008 | TORCH_CHECK( |
| 1009 | hc_scale.scalar_type() == at::kFloat && hc_scale.numel() == 1, |
| 1010 | "hc_head_fused_cpu: hc_scale must be a scalar float32"); |
| 1011 | TORCH_CHECK(hc_base.scalar_type() == at::kFloat, "hc_head_fused_cpu: hc_base must be float32"); |
| 1012 | TORCH_CHECK(hc_base.numel() == hc_mult, "hc_head_fused_cpu: hc_base.numel() must equal hc_mult"); |
| 1013 | TORCH_CHECK( |
| 1014 | x.device().is_cpu() && hc_fn.device().is_cpu() && hc_scale.device().is_cpu() && hc_base.device().is_cpu(), |
| 1015 | "hc_head_fused_cpu: all inputs must be CPU tensors"); |
| 1016 | |
| 1017 | const int64_t T = x.size(0); |
| 1018 | const int64_t d = x.size(2); |
| 1019 | TORCH_CHECK(T > 0 && d > 0, "hc_head_fused_cpu: T and d must be positive"); |
| 1020 | TORCH_CHECK( |
| 1021 | hc_fn.dim() == 2 && hc_fn.size(0) == hc_mult && hc_fn.size(1) == hc_mult * d, |
| 1022 | "hc_head_fused_cpu: hc_fn must be [hc_mult, hc_mult*d]"); |
| 1023 | |
| 1024 | TORCH_CHECK( |
| 1025 | x.is_contiguous() && hc_fn.is_contiguous() && hc_scale.is_contiguous() && hc_base.is_contiguous(), |
| 1026 | "hc_head_fused_cpu: x/hc_fn/hc_scale/hc_base must be contiguous"); |
| 1027 | const float hc_scale_val = hc_scale.data_ptr<float>()[0]; |
| 1028 | |
| 1029 | auto y = at::empty({T, d}, x.options()); |
| 1030 | |
| 1031 | AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_head_fused_cpu", [&] { |
| 1032 | hc_head_fuse_impl<scalar_t>( |
| 1033 | y.data_ptr<scalar_t>(), |
| 1034 | x.data_ptr<scalar_t>(), |
| 1035 | hc_fn.data_ptr<float>(), |
| 1036 | hc_scale_val, |
| 1037 | hc_base.data_ptr<float>(), |
| 1038 | T, |
| 1039 | d, |
| 1040 | static_cast<int>(hc_mult), |
| 1041 | static_cast<float>(hc_eps), |
| 1042 | static_cast<float>(norm_eps)); |
| 1043 | }); |
| 1044 | |
| 1045 | return y; |
| 1046 | } |