[KernelSwift算子优化] 程柯雷-Task03 norm_fn-571900729

dlblas/kernels/ks_competition/QY2026_Autumn_Lfan-ke_Task03NormFn_MetaXC500.py

@@ -0,0 +1,260 @@
+"""Task03 norm_fn — 优化实现 (v1)：一个融合 MXMACA kernel 顶掉整段 torch 参考。
+
+参考实现把这件事摊成了一串独立算子：bf16→fp32 的 cast、einsum 走的那次小 GEMM、square、
+sum、rsqrt、逐元素乘、再 sum、最后 view——每一步都是一次 kernel 启动加一趟中间张量往返显存，
+而真正要算的只有 13 行 × 24 个长度 5120 的点积，外加每行一个平方和。1.6M 次乘加对 C500 来说
+是零头，所以这个算子跟 engram_hash 一样：瓶颈在启动与访存，优化就是压成一次 launch。
+
+kernel 布局：一个 block 负责一个 (row, j) 输出元素，共 13×24 = 312 个 block。每个 block 只做
+一趟读、一次块内归约：沿 K 走的时候同时累 `acc += r*f`（点积）与 `sq += r*r`（RMS 分母），
+读 residual 与 fn[j] 都是顺序合并访存，末了 0 号线程写 `dot * rsqrt(sq/K + eps)`。
+
+★这个布局是量出来换的，不是一开始就这么写的。第一版是「一 block 一行、块内串 24 次归约」：
+正确，但只有 13 个 block、每次归约都要 `__syncthreads`，实测 0.357 ms 对 v0 的 0.234 ms——
+比参考还慢 (0.66×)。改成 312 个 block、每块一次归约后才跑赢。sqrsum 在 24 个 block 间重复算，
+多读 23 遍 residual（10 KB/行）换掉全部同步，这笔账在这个尺寸上是划算的。
+
+★harness 约束：auto_bench.py 的 `_filter_module_ast` 只保留 import、class/def 和字面量赋值，
+模块级的其它语句会被静默丢弃，所以扩展必须在函数里懒加载。
+"""
+
+import os
+
+import torch
+from torch.utils.cpp_extension import load_inline
+
+n1 = 13
+mhc_mult = 4
+hidden_size = 1280
+generate_normw = False
+
+# C500 上扫过 128/256/512/1024，见 README 的调参记录。
+DEFAULT_BLOCK = 256
+
+
+def _cuda_source():
+    return r"""
+#include <torch/extension.h>
+#include <cstdint>
+
+// MACA 的 bf16 叫 __maca_bfloat16 且没有到 float 的转换算子，所以干脆不碰那个类型：
+// bf16 转 fp32 就是把这 16 位放到 fp32 的高 16 位，低位补零——位级精确，也不依赖任何内建函数。
+__device__ __forceinline__ float bf16_to_f32(uint16_t b) {
+  union { uint32_t u; float f; } c;
+  c.u = (uint32_t)b << 16;
+  return c.f;
+}
+
+// 一 block 负责同一行的 JG 个 j。residual 与 RMS 分母整块只过一遍、JG 个点积在寄存器里并行累，
+// 把参考布局里「每个 j 都把 residual 重读一遍、平方和重算一遍」的 24 倍冗余摊掉。
+// JG 是量出来的：1/2/3/4/6/8/12 对应 312/156/104/78/52/39/26 个 block，
+// 实测 37.79/36.86/36.10/35.85/35.27/37.07/42.82 us —— 6 是甜点。再大 block 数不够，占用率掉下来。
+template <int JG>
+__global__ void norm_fn_kernel(
+    const uint16_t* __restrict__ residual,   // [rows, K] bf16 的原始位
+    const float* __restrict__ fn,            // [M, K] fp32
+    float* __restrict__ out,                 // [rows, M] fp32
+    int M, int K, float eps, float inv_K) {
+  extern __shared__ float red[];             // (JG+1) * nthreads：JG 个点积 + 一个平方和
+
+  const int ng = M / JG;
+  const int row = blockIdx.x / ng;
+  const int j0 = (blockIdx.x - row * ng) * JG;
+  const int tid = threadIdx.x;
+  const int nt = blockDim.x;
+
+  // K 是 4 的倍数（本题 5120），一次取 4 个：fn 走 float4，residual 走 4 个 bf16 打包成的 uint2。
+  const uint2* r4 = reinterpret_cast<const uint2*>(residual + (long)row * K);
+  const float4* f4[JG];
+#pragma unroll
+  for (int u = 0; u < JG; ++u) f4[u] = reinterpret_cast<const float4*>(fn + (long)(j0 + u) * K);
+
+  float acc[JG];
+#pragma unroll
+  for (int u = 0; u < JG; ++u) acc[u] = 0.f;
+  float sq = 0.f;
+
+  const int K4 = K >> 2;
+  for (int k = tid; k < K4; k += nt) {
+    const uint2 rb = r4[k];
+    const float r0 = bf16_to_f32((uint16_t)(rb.x & 0xffffu));
+    const float r1 = bf16_to_f32((uint16_t)(rb.x >> 16));
+    const float r2 = bf16_to_f32((uint16_t)(rb.y & 0xffffu));
+    const float r3 = bf16_to_f32((uint16_t)(rb.y >> 16));
+    sq = fmaf(r0, r0, sq); sq = fmaf(r1, r1, sq);
+    sq = fmaf(r2, r2, sq); sq = fmaf(r3, r3, sq);
+#pragma unroll
+    for (int u = 0; u < JG; ++u) {
+      const float4 fv = f4[u][k];
+      acc[u] = fmaf(r0, fv.x, acc[u]);
+      acc[u] = fmaf(r1, fv.y, acc[u]);
+      acc[u] = fmaf(r2, fv.z, acc[u]);
+      acc[u] = fmaf(r3, fv.w, acc[u]);
+    }
+  }
+
+#pragma unroll
+  for (int u = 0; u < JG; ++u) red[u * nt + tid] = acc[u];
+  red[JG * nt + tid] = sq;
+  __syncthreads();
+  for (int s = nt >> 1; s > 0; s >>= 1) {
+    if (tid < s) {
+#pragma unroll
+      for (int u = 0; u <= JG; ++u) red[u * nt + tid] += red[u * nt + tid + s];
+    }
+    __syncthreads();
+  }
+  if (tid < JG) out[(long)row * M + j0 + tid] = red[tid * nt] * rsqrtf(red[JG * nt] * inv_K + eps);
+}
+
+torch::Tensor norm_fn(torch::Tensor residual_, torch::Tensor fn_, double eps,
+                      int64_t block) {
+  auto residual = residual_.is_contiguous() ? residual_ : residual_.contiguous();
+  auto fn = fn_.is_contiguous() ? fn_ : fn_.contiguous();
+
+  // residual: [n0, n1, mhc_mult, hidden]  ->  [n0*n1, K]
+  const int n0 = (int)residual.size(0);
+  const int n1v = (int)residual.size(1);
+  const int rows = n0 * n1v;
+  const int M = (int)fn.size(0);
+  const int K = (int)fn.size(1);
+
+  TORCH_CHECK((K & 3) == 0, "K must be a multiple of 4");
+
+  auto out = torch::empty({rows, M}, fn.options());
+  const int threads = (int)block;
+
+  static const int kCand[] = {6, 4, 3, 2};
+  int JG = 1;
+  for (int i = 0; i < 4; ++i) {
+    if (M % kCand[i] == 0) { JG = kCand[i]; break; }
+  }
+  const size_t shmem = (size_t)(JG + 1) * threads * sizeof(float);
+  const int grid = rows * (M / JG);
+  auto pr = (const uint16_t*)residual.data_ptr<at::BFloat16>();
+  auto pf = fn.data_ptr<float>();
+  auto po = out.data_ptr<float>();
+  const float e = (float)eps, ik = 1.0f / (float)K;
+  switch (JG) {
+    case 6: norm_fn_kernel<6><<<grid, threads, shmem>>>(pr, pf, po, M, K, e, ik); break;
+    case 4: norm_fn_kernel<4><<<grid, threads, shmem>>>(pr, pf, po, M, K, e, ik); break;
+    case 3: norm_fn_kernel<3><<<grid, threads, shmem>>>(pr, pf, po, M, K, e, ik); break;
+    case 2: norm_fn_kernel<2><<<grid, threads, shmem>>>(pr, pf, po, M, K, e, ik); break;
+    default: norm_fn_kernel<1><<<grid, threads, shmem>>>(pr, pf, po, M, K, e, ik); break;
+  }
+  return out.view({n0, n1v, M});
+}
+"""
+
+
+def _mod(_cache=[]):
+    """编译一次、缓存住的扩展。放在函数里是 harness 的硬要求（见模块头）。"""
+    if not _cache:
+        _cache.append(
+            load_inline(
+                name="norm_fn_maca",
+                cpp_sources=(
+                    "torch::Tensor norm_fn(torch::Tensor residual, torch::Tensor fn, "
+                    "double eps, int64_t block);"
+                ),
+                cuda_sources=_cuda_source(),
+                functions=["norm_fn"],
+                # ★宿主侧（pybind 胶水 + main.cpp）默认是 -O0：torch 不注入 -O、cu-bridge 也不注入，
+                # 而 pybind 往返占这台机器 26 µs 硬底里的 8.9 µs。extra_cuda_cflags 加 -O3 是空操作
+                # （设备段 md5 逐字节相同，已证），真正的口子在 extra_cflags。
+                # 官方 harness 交替 A/B 五对全胜：5.475/5.473/5.477/5.430/5.511 -> 5.666/5.741/5.774/5.868/6.125
+                extra_cflags=["-O3"],
+                verbose=False,
+            )
+        )
+    return _cache[0]
+
+
+def _block_size():
+    return int(os.environ.get("NORM_FN_BLOCK", DEFAULT_BLOCK))
+
+
+class Model(torch.nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(
+        self,
+        residual: torch.Tensor,
+        mhc_fn: torch.Tensor,
+        mhc_norm_weight: torch.Tensor | None,
+        mhc_norm_eps: float,
+    ) -> torch.Tensor:
+        if mhc_norm_weight is not None:
+            mhc_fn = mhc_fn * mhc_norm_weight
+        residual = residual.flatten(2, 3).float()
+        assert mhc_fn.dtype == residual.dtype == torch.float
+        mhc_mult = mhc_fn.shape[0]
+        rms_group_size = mhc_fn.shape[-1]
+        mixes = torch.einsum(
+            "mbk,nbk->mbn",
+            residual.view(-1, 1, rms_group_size),
+            mhc_fn.view(mhc_mult, 1, rms_group_size),
+        )
+        sqrsum = residual.view(-1, 1, rms_group_size).square().sum(-1)
+        mixes = (
+            mixes * (sqrsum.unsqueeze(-1) / rms_group_size + mhc_norm_eps).rsqrt()
+        ).sum(-2)
+        return mixes.view(*residual.shape[:2], -1)
+
+
+class ModelNew(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._op = _mod().norm_fn
+        self._block = _block_size()
+
+    def forward(
+        self,
+        residual: torch.Tensor,
+        mhc_fn: torch.Tensor,
+        mhc_norm_weight: torch.Tensor | None,
+        mhc_norm_eps: float,
+    ) -> torch.Tensor:
+        if mhc_norm_weight is not None:
+            mhc_fn = mhc_fn * mhc_norm_weight
+        return self._op(residual, mhc_fn, mhc_norm_eps, self._block)
+
+
+def generate_norm_fn_test_data(n1, mhc_mult, hidden_size, generate_normw, device="cuda"):
+    n0 = 1
+    mhc_mult3 = mhc_mult * (2 + mhc_mult)
+    mhc_hidden_size = mhc_mult * hidden_size
+
+    residual = (
+        torch.randn((n0, n1, mhc_mult, hidden_size), dtype=torch.float, device=device)
+        .mul(1 + torch.arange(mhc_mult, device=device).mul(0.01).view(1, 1, -1, 1))
+        .bfloat16()
+    )
+
+    fn = (
+        torch.randn((mhc_mult3, mhc_mult, hidden_size), dtype=torch.float, device=device)
+        * 1e-4
+        * (1 + torch.arange(mhc_mult, device=device).mul(0.01).view(1, -1, 1))
+    ).flatten(1, 2)
+
+    if generate_normw:
+        normw = torch.randn((mhc_hidden_size,), dtype=torch.float, device=device) * 0.1 + 1.0
+    else:
+        normw = None
+
+    out_grad = torch.randn((n0, n1, mhc_mult3), dtype=torch.float, device=device)
+
+    return [residual, fn, normw, out_grad, 1e-6]
+
+
+def get_inputs():
+    torch.manual_seed(233)
+    residual, fn, normw, out_grad, mhc_norm_eps = generate_norm_fn_test_data(
+        n1, mhc_mult, hidden_size, generate_normw
+    )
+    return [residual, fn, None, mhc_norm_eps]
+
+
+def get_init_inputs():
+    return []