diff --git a/dlblas/kernels/ks_competition/QY2026_Autumn_Lfan-ke_Task03NormFn_MetaXC500.py b/dlblas/kernels/ks_competition/QY2026_Autumn_Lfan-ke_Task03NormFn_MetaXC500.py new file mode 100644 index 00000000..02df4392 --- /dev/null +++ b/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 +#include + +// 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 +__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(residual + (long)row * K); + const float4* f4[JG]; +#pragma unroll + for (int u = 0; u < JG; ++u) f4[u] = reinterpret_cast(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(); + auto pf = fn.data_ptr(); + auto po = out.data_ptr(); + const float e = (float)eps, ik = 1.0f / (float)K; + switch (JG) { + case 6: norm_fn_kernel<6><<>>(pr, pf, po, M, K, e, ik); break; + case 4: norm_fn_kernel<4><<>>(pr, pf, po, M, K, e, ik); break; + case 3: norm_fn_kernel<3><<>>(pr, pf, po, M, K, e, ik); break; + case 2: norm_fn_kernel<2><<>>(pr, pf, po, M, K, e, ik); break; + default: norm_fn_kernel<1><<>>(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 []