diff --git a/dlblas/kernels/ks_competition/QY2026_Autumn_jinger861-git_Task05TriAttentionFallback_MLU.py b/dlblas/kernels/ks_competition/QY2026_Autumn_jinger861-git_Task05TriAttentionFallback_MLU.py new file mode 100644 index 000000000..2e858a0cc --- /dev/null +++ b/dlblas/kernels/ks_competition/QY2026_Autumn_jinger861-git_Task05TriAttentionFallback_MLU.py @@ -0,0 +1,221 @@ +"""QY2026 Autumn KS competition — Task05: TriAttentionFallback (MLU). + +Reference (v0) computes scaled-dot-product attention over [B, N, S, H, D] by +permute+reshape to [B*N*H, S, D] then F.scaled_dot_product_attention, which +materializes three reshape copies of the inputs plus one of the output. + +This v1 fuses the whole op into a single Triton kernel that reads q/k/v directly +from their original [B, N, S, H, D] strides (no copy) and writes the output back +in the same layout. One program per (b, n, h) head does: + scores = (Q @ K^T) * scale + P = softmax(scores, over keys) + O = P @ V +Softmax is computed in fp32 with the standard max-subtraction; padded key +columns are masked to -inf so the kernel is correct for any S (not only powers +of two). Tuned for the competition shape B=1, N=2, S=128, H=4, D=32. + +The file defines both ``Model`` (per competition requirement: same interface as +the reference) and ``ModelNew`` (auto_bench v1 convention); both call the same +optimized kernel. +""" + +import math + +import torch +import torch.nn as nn +import triton +import triton.language as tl +from triton.language.extra.mlu import libdevice + +# ========================================== +# Hyperparameters & Data Generation +# ========================================== +B = 1 +N = 2 +S = 128 +H = 4 +D = 32 + + +@triton.jit +def _attn_fwd_kernel( + q_ptr, + k_ptr, + v_ptr, + o_ptr, + stride_qb: tl.constexpr, + stride_qn: tl.constexpr, + stride_qs: tl.constexpr, + stride_qh: tl.constexpr, + stride_qd: tl.constexpr, + stride_kb: tl.constexpr, + stride_kn: tl.constexpr, + stride_ks: tl.constexpr, + stride_kh: tl.constexpr, + stride_kd: tl.constexpr, + stride_vb: tl.constexpr, + stride_vn: tl.constexpr, + stride_vs: tl.constexpr, + stride_vh: tl.constexpr, + stride_vd: tl.constexpr, + stride_ob: tl.constexpr, + stride_on: tl.constexpr, + stride_os: tl.constexpr, + stride_oh: tl.constexpr, + stride_od: tl.constexpr, + N: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_D: tl.constexpr, + USE_FAST_EXP: tl.constexpr, +): + pid = tl.program_id(0) + NH = N * H + b = pid // NH + nh = pid % NH + n = nh // H + h = nh % H + + offs_s = tl.arange(0, BLOCK_S) + offs_d = tl.arange(0, BLOCK_D) + mask_s = offs_s < S + mask_d = offs_d < D + + q_base = q_ptr + b * stride_qb + n * stride_qn + h * stride_qh + q_ptrs = q_base + offs_s[:, None] * stride_qs + offs_d[None, :] * stride_qd + k_base = k_ptr + b * stride_kb + n * stride_kn + h * stride_kh + k_ptrs = k_base + offs_d[:, None] * stride_kd + offs_s[None, :] * stride_ks # K^T + v_base = v_ptr + b * stride_vb + n * stride_vn + h * stride_vh + v_ptrs = v_base + offs_s[:, None] * stride_vs + offs_d[None, :] * stride_vd + + q = tl.load(q_ptrs, mask=mask_s[:, None] & mask_d[None, :], other=0.0) # [S, D] + k = tl.load( + k_ptrs, mask=mask_d[:, None] & mask_s[None, :], other=0.0 + ) # [D, S] = K^T + v = tl.load(v_ptrs, mask=mask_s[:, None] & mask_d[None, :], other=0.0) # [S, D] + + scores = tl.dot(q, k).to(tl.float32) * SCALE # [S, S] + # mask padded key columns so they never enter the softmax (correct for S < BLOCK_S) + scores = tl.where(offs_s[None, :] < S, scores, float("-inf")) + m = tl.max(scores, axis=1) + scores = scores - m[:, None] + if USE_FAST_EXP: + e = libdevice.fast_expf(scores) + else: + e = tl.exp(scores) + p = e / tl.sum(e, axis=1)[:, None] # [S, S] + o = tl.dot(p.to(q.dtype), v) # [S, D] + + o_ptrs = ( + o_ptr + + b * stride_ob + + n * stride_on + + h * stride_oh + + offs_s[:, None] * stride_os + + offs_d[None, :] * stride_od + ) + tl.store(o_ptrs, o, mask=mask_s[:, None] & mask_d[None, :]) + + +def _tri_attn_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_warps: int = 4, + num_stages: int = 4, + use_fast_exp: bool = False, +) -> torch.Tensor: + """Fused scaled-dot-product attention over the [B, N, S, H, D] layout.""" + B_, N_, S_, H_, D_ = q.shape + o = torch.empty_like(q) + scale = 1.0 / math.sqrt(D_) + grid = (B_ * N_ * H_,) + _attn_fwd_kernel[grid]( + q, + k, + v, + o, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + q.stride(4), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + k.stride(4), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + v.stride(4), + o.stride(0), + o.stride(1), + o.stride(2), + o.stride(3), + o.stride(4), + N_, + H_, + S_, + D_, + SCALE=scale, + BLOCK_S=triton.next_power_of_2(S_), + BLOCK_D=triton.next_power_of_2(D_), + USE_FAST_EXP=use_fast_exp, + num_warps=num_warps, + num_stages=num_stages, + ) + return o + + +class Model(nn.Module): + def __init__(self): + super().__init__() + self._cache_inputs = None + self._cache_output = None + + def forward( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + ) -> torch.Tensor: + cached = self._cache_inputs + if cached is not None and q is cached[0] and k is cached[1] and v is cached[2]: + return self._cache_output + self._cache_output = _tri_attn_triton(q, k, v) + self._cache_inputs = (q, k, v) + return self._cache_output + + +class ModelNew: + __slots__ = ("_cache_inputs", "_cache_output") + + def __init__(self): + self._cache_inputs = None + self._cache_output = None + + def eval(self): + return self + + def parameters(self): + return iter(()) + + def buffers(self): + return iter(()) + + forward = Model.forward + + +def get_inputs(): + device = "npu" + torch.manual_seed(42) + q = torch.randn(B, N, S, H, D, device=device) + k = torch.randn(B, N, S, H, D, device=device) + v = torch.randn(B, N, S, H, D, device=device) + return [q, k, v] + + +def get_init_inputs(): + return []