diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head.py new file mode 100644 index 0000000..41d15f3 --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head.py @@ -0,0 +1,112 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 国际A(NVIDIA)/国际B(AMD)通用版 · v2 +DeepSeek-V4 "hc_head" LM-head 混合器: RMSNorm + 线性混合 + sigmoid 门控 + 加权求和 +v2 修复(从 extend_attention.py v2 经验): + 1) num_stages=1 显式 (头号嫌疑 — 国产 NPU + 部分国际卡多级流水支持不完整) + 2) num_warps 4 → 8 (大 BLOCK_D=1024 + HC_MULT=4 循环, 8 warps 更稳) + 3) pid int64 → int32 (D ≤ 7168, T 不会超 2.1B, 长度算术走 int32) + 4) 指针 cast 防 NaN (D245 验证) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + # 指针 cast 防 NaN(D245 验证) + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + # ===== Pass 1: 算 squared sum + 算 mixes (linear projection) ===== + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + # ===== sigmoid gate ===== + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps # [HC_MULT] + + # ===== Pass 2: 加权求和 (collapse hc_mult) ===== + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + """hc_head: DeepSeek-V4 HC head reduction for LM-head mixer. + + Computes gates from the RMS-normalized flattened HC residual + and returns out = sum_i gate_i * residual_i, collapsing hc_mult streams. + + Args: + x: [T, hc_mult, hidden_size] bfloat16 + hc_fn: [hc_mult, hc_mult * hidden_size] float32 + hc_scale: [1] float32 + hc_base: [hc_mult] float32 + norm_eps, hc_eps: float + + Returns: + [T, hidden_size] bfloat16 + """ + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=8, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_ascend.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_ascend.py new file mode 100644 index 0000000..6c5dc53 --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_ascend.py @@ -0,0 +1,93 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 华为昇腾(Ascend)特化版 · v2 +算法层同通用版 · 华为国产 NPU 套路: + 1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需) + 2) enable_fp_fusion=False (D245 验证) + 3) num_stages=1 (国产 NPU 多级流水支持不完整) + 4) num_warps=8 (大 BLOCK_D=1024 + HC_MULT=4 循环) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps + + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + with torch.get_device_module(x.device).device(x.device): + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=8, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_enflame.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_enflame.py new file mode 100644 index 0000000..9c1f94d --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_enflame.py @@ -0,0 +1,95 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 燧原(Enflame)特化版 · v2 +算法层同通用版 · 燧原小芯片套路(参考 chunk_scaled_dot_kkt_enflame.py 模式): + 1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需) + 2) enable_fp_fusion=False (D245 验证) + 3) BLOCK_D: 1024 → 512 (小芯片 + 小 SRAM, 跟 chunk_enflame BK=32 一致思路) + 4) num_warps: 8 → 4 (小芯片 warp 资源少) + 5) num_stages=1 (国产 NPU 多级流水支持不完整) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps + + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + # 燧原小 SRAM → 小 BLOCK + BLOCK_D = min(512, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + with torch.get_device_module(x.device).device(x.device): + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=4, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_hygon.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_hygon.py new file mode 100644 index 0000000..49ce13d --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_hygon.py @@ -0,0 +1,93 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 海光(Hygon DCU)特化版 · v2 +算法层同通用版 · 海光 DCU 套路: + 1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需) + 2) enable_fp_fusion=False (D245 验证) + 3) num_stages=1 (DCU 多级流水支持不完整) + 4) num_warps=8 (大 BLOCK_D=1024) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps + + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + with torch.get_device_module(x.device).device(x.device): + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=8, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_iluvatar.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_iluvatar.py new file mode 100644 index 0000000..e26a893 --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_iluvatar.py @@ -0,0 +1,93 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 天数智芯(Iluvatar)特化版 · v2 +算法层同通用版 · 天数智芯套路: + 1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需) + 2) enable_fp_fusion=False (D245 验证) + 3) num_stages=1 (国产 NPU 多级流水支持不完整) + 4) num_warps=8 (大 BLOCK_D=1024) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps + + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + with torch.get_device_module(x.device).device(x.device): + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=8, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_kunlunxin.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_kunlunxin.py new file mode 100644 index 0000000..98245d3 --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_kunlunxin.py @@ -0,0 +1,96 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 昆仑芯(Kunlunxin)特化版 · v2 +算法层同通用版 · 昆仑芯套路(参考 chunk_scaled_dot_kkt_kunlunxin.py 模式): + 1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需) + 2) enable_fp_fusion=False (D245 验证) + 3) num_stages=1 (SDNN 路径多级流水 codegen bug — 强制单级) + 4) num_warps=8 (大 BLOCK_D=1024) + 5) tl.static_range 已用 (强制编译时展开 HC_MULT 循环, 绕 SDNN 循环 codegen bug) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + # tl.static_range 强制编译时展开(绕 SDNN 循环 codegen bug) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps + + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + # tl.static_range 强制编译时展开(绕 SDNN 循环 codegen bug) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + with torch.get_device_module(x.device).device(x.device): + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=8, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_metax.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_metax.py new file mode 100644 index 0000000..18b79c5 --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/hc_head_metax.py @@ -0,0 +1,93 @@ +"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 沐曦(Metax)特化版 · v2 +算法层同通用版 · 沐曦套路(参考 chunk_scaled_dot_kkt_metax.py 模式): + 1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需) + 2) enable_fp_fusion=False (D245 验证) + 3) num_stages=1 (国产 NPU 多级流水支持不完整, 不写用默认也安全 — 显式写更稳) + 4) num_warps=8 (大 BLOCK_D=1024) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_head_fwd_kernel( + x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr, + T, + HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr, + BLOCK_D: tl.constexpr, + norm_eps, hc_eps, +): + pid = tl.program_id(0).to(tl.int32) + + if x_ptr.dtype.element_ty.primitive_bitwidth == 16: + x_ptr = x_ptr.to(tl.pointer_type(tl.int16)) + if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32: + hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32)) + + sqr_sum = tl.zeros((), dtype=tl.float32) + mixes = tl.zeros((HC_MULT,), dtype=tl.float32) + + for d_off in range(0, HC_DIM, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < HC_DIM + x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32) + sqr_sum += tl.sum(x_val * x_val, axis=0) + for m in tl.static_range(HC_MULT): + fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0) + mixes[m] += tl.sum(x_val * fn_val, axis=0) + + rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps) + mixes = mixes * rsqrt + + scale = tl.load(hc_scale_ptr) + bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT)) + pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps + + for d_off in range(0, D, BLOCK_D): + d_idx = d_off + tl.arange(0, BLOCK_D) + mask = d_idx < D + accum = tl.zeros((BLOCK_D,), dtype=tl.float32) + for m in tl.static_range(HC_MULT): + x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32) + accum += pre[m] * x_val + tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask) + + +def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + if x.ndim != 3: + raise ValueError('x must have shape [T, hc_mult, hidden_size]') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + + if hc_fn.shape != (hc_mult, HC_DIM): + raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}') + if hc_scale.numel() != 1: + raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}') + if hc_base.shape != (hc_mult,): + raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}') + + x_flat = x.view(T, HC_DIM) + out = torch.empty((T, D), dtype=x.dtype, device=x.device) + if T == 0: + return out + + BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length())) + + grid = (T,) + with torch.get_device_module(x.device).device(x.device): + _hc_head_fwd_kernel[grid]( + x_flat, hc_fn, hc_scale, hc_base, out, + T, + HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult, + BLOCK_D=BLOCK_D, + norm_eps=norm_eps, hc_eps=hc_eps, + num_warps=8, num_stages=1, enable_fp_fusion=False, + ) + return out + + +reference = hc_head diff --git a/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/test_8files_algorithm.py b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/test_8files_algorithm.py new file mode 100644 index 0000000..87600ff --- /dev/null +++ b/eternal-lake-heart/love-core/see-you-tomorrow-channel/zz-day-5/results/Task55_hc_head_8files/test_8files_algorithm.py @@ -0,0 +1,111 @@ +"""Task 55 hc_head v1 算法层测试 + +用 torch 模拟 v1 kernel 的逻辑, 跟平台给的 reference 比较 +(bit-exact 不可能, bf16 有累积误差, 用 atol=1.5e-2, rtol=1.5e-2 平台标准) + +注意: 这个测试是 torch 模拟 v1 算法, 验证"如果 v1 在真 GPU 上跑, 算法层是否对" +""" +import torch +import torch.nn.functional as F +import sys +sys.path.insert(0, '.') + + +# ============================================================ +# 平台给的 reference (照抄) +# ============================================================ +def reference(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + shape, dtype = x.size(), x.dtype + x = x.flatten(1).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1) + return y.to(dtype) + + +# ============================================================ +# v1 的 torch 模拟(模拟 K1 的 2-pass 逻辑) +# ============================================================ +def v1_torch(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps): + """v1 的 torch 模拟: 跟 K1 逻辑 bit-exact""" + T, hc_mult, D = x.shape + HC_DIM = hc_mult * D + x_flat = x.view(T, HC_DIM).float() + hc_fn_f = hc_fn.float() + scale = hc_scale.float().item() + bases = hc_base.float() + + # Pass 1: sqr_sum + mixes + sqr_sum = x_flat.square().sum(-1) # [T] + rsqrt = torch.rsqrt(sqr_sum / HC_DIM + norm_eps) # [T] + mixes = F.linear(x_flat, hc_fn_f) * rsqrt.unsqueeze(-1) # [T, HC_MULT] + + # sigmoid gate + pre = torch.sigmoid(mixes * scale + bases) + hc_eps # [T, HC_MULT] + + # Pass 2: 加权求和 + x_back = x.view(T, hc_mult, D).float() # [T, HC_MULT, D] + y = (pre.unsqueeze(-1) * x_back).sum(dim=1) # [T, D] + return y.to(x.dtype) + + +# ============================================================ +# 跑测试 +# ============================================================ +def test_case(T, hc_mult, D, norm_eps=1e-6, hc_eps=1e-6, seed=42): + """跑一个测试用例, 比较 reference 和 v1_torch""" + torch.manual_seed(seed) + x = torch.randn(T, hc_mult, D, dtype=torch.bfloat16) * 0.5 + hc_fn = torch.randn(hc_mult, hc_mult * D, dtype=torch.float32) * 0.02 + hc_scale = torch.tensor([2.0], dtype=torch.float32) + hc_base = torch.randn(hc_mult, dtype=torch.float32) * 0.1 + + ref_out = reference(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps) + v1_out = v1_torch(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps) + + # 平台标准: atol=1.5e-2, rtol=1.5e-2 + match = torch.allclose(v1_out.float(), ref_out.float(), atol=1.5e-2, rtol=1.5e-2) + max_diff = (v1_out.float() - ref_out.float()).abs().max().item() + max_ref = ref_out.float().abs().max().item() + rel_diff = max_diff / max(max_ref, 1e-9) + + print(f" T={T:5d} hc_mult={hc_mult} D={D:5d} " + f"match={match} max_abs_diff={max_diff:.6f} max_rel_diff={rel_diff:.6f}") + return match + + +if __name__ == "__main__": + print("Testing Task 55 hc_head v1 vs reference algorithm equivalence...") + print("(平台标准: atol=1.5e-2, rtol=1.5e-2)") + print() + + test_cases = [ + # (T, hc_mult, D) 默认 hc_mult=4 (DeepSeek-V4), D=7168 + (1, 4, 7168), # 1 token + (8, 4, 7168), # 8 token + (32, 4, 7168), # 32 token + (128, 4, 7168), # 128 token (典型 batch) + (1, 4, 1024), # 1 token 小 D + (8, 4, 256), # 8 token 更小 D + (16, 2, 128), # hc_mult=2, 小 D + (1, 4, 128), # 极小 + ] + + all_pass = True + for T, hc_mult, D in test_cases: + ok = test_case(T, hc_mult, D) + all_pass = all_pass and ok + + # 再用不同 norm_eps / hc_eps 测 + print() + print("Extra edge cases (不同 eps):") + for eps in [1e-8, 1e-4, 1e-2]: + ok = test_case(8, 4, 1024, norm_eps=eps, hc_eps=eps) + all_pass = all_pass and ok + + print() + print("=" * 70) + print(f"{'ALL PASS ✓' if all_pass else 'SOME FAILED ✗'}") + print("=" * 70) + sys.exit(0 if all_pass else 1)