[2026秋季][Task04] Foxxiaohuli

dlblas/kernels/ks_competition/triton/QY2026_Autumn_Foxxiaohuli_Task04RelativePositionEncoding_MetaX.py

@@ -0,0 +1,343 @@
+TILE = 16
+ALIGN = 8
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import triton
+import triton.language as tl
+
+
+def generate_relp(
+    *,
+    asym_id: torch.Tensor,
+    residue_index: torch.Tensor,
+    entity_id: torch.Tensor,
+    token_index: torch.Tensor,
+    sym_id: torch.Tensor,
+    r_max: int = 32,
+    s_max: int = 2,
+) -> torch.Tensor:
+
+    b_same_chain = (asym_id[:, None] == asym_id[None, :]).long()
+    b_same_residue = (residue_index[:, None] == residue_index[None, :]).long()
+    b_same_entity = (entity_id[:, None] == entity_id[None, :]).long()
+
+    d_residue = torch.clamp(
+        residue_index[:, None] - residue_index[None, :] + r_max, min=0, max=2 * r_max
+    )
+
+    d_residue = d_residue * b_same_chain + (1 - b_same_chain) * (2 * r_max + 1)
+    a_rel_pos = F.one_hot(d_residue, 2 * (r_max + 1))
+
+    d_token = torch.clamp(
+        token_index[:, None] - token_index[None, :] + r_max, min=0, max=2 * r_max
+    )
+
+    d_token = d_token * b_same_chain * b_same_residue + (
+        1 - b_same_chain * b_same_residue
+    ) * (2 * r_max + 1)
+    a_rel_token = F.one_hot(d_token, 2 * (r_max + 1))
+
+    d_chain = torch.clamp(
+        sym_id[:, None] - sym_id[None, :] + s_max, min=0, max=2 * s_max
+    )
+
+    d_chain = d_chain * b_same_entity + (1 - b_same_entity) * (2 * s_max + 1)
+    a_rel_chain = F.one_hot(d_chain, 2 * (s_max + 1))
+
+    relp = torch.cat(
+        [
+            a_rel_pos,
+            a_rel_token,
+            b_same_entity[..., None],
+            a_rel_chain,
+        ],
+        dim=-1,
+    ).float()
+    return relp
+
+
+@triton.jit
+def _fused_relpos_proj_kernel(
+    asym_id_ptr,
+    res_idx_ptr,
+    ent_id_ptr,
+    tok_idx_ptr,
+    sym_id_ptr,
+    W_T_ptr,
+    W_T_stride0,
+    out_ptr,
+    out_s0,
+    out_s1,
+    out_s2,
+    N,
+    C_Z,
+    r_max,
+    s_max,
+    off_res,
+    off_tok,
+    off_ent,
+    off_chain,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    m_offsets = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    n_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    m_mask = m_offsets < N
+    n_mask = n_offsets < N
+
+    a_i = tl.load(asym_id_ptr + m_offsets, mask=m_mask, other=0).to(tl.int32)
+    a_j = tl.load(asym_id_ptr + n_offsets, mask=n_mask, other=1).to(tl.int32)
+
+    ri = tl.load(res_idx_ptr + m_offsets, mask=m_mask, other=0).to(tl.int32)
+    rj = tl.load(res_idx_ptr + n_offsets, mask=n_mask, other=0).to(tl.int32)
+
+    ei = tl.load(ent_id_ptr + m_offsets, mask=m_mask, other=0).to(tl.int32)
+    ej = tl.load(ent_id_ptr + n_offsets, mask=n_mask, other=1).to(tl.int32)
+
+    ti = tl.load(tok_idx_ptr + m_offsets, mask=m_mask, other=0).to(tl.int32)
+    tj = tl.load(tok_idx_ptr + n_offsets, mask=n_mask, other=0).to(tl.int32)
+
+    si = tl.load(sym_id_ptr + m_offsets, mask=m_mask, other=0).to(tl.int32)
+    sj = tl.load(sym_id_ptr + n_offsets, mask=n_mask, other=0).to(tl.int32)
+
+    same_chain = a_i[:, None] == a_j[None, :]
+    same_residue = ri[:, None] == rj[None, :]
+    same_entity = ei[:, None] == ej[None, :]
+    same_chain_res = same_chain & same_residue
+
+    pair_valid = m_mask[:, None] & n_mask[None, :]
+
+    two_r = 2 * r_max
+    d_res_raw = ri[:, None] - rj[None, :] + r_max
+    d_res_clamped = tl.where(d_res_raw < 0, 0, d_res_raw)
+    d_res_clamped = tl.where(d_res_clamped > two_r, two_r, d_res_clamped)
+    d_res = tl.where(same_chain, d_res_clamped, two_r + 1)
+
+    d_tok_raw = ti[:, None] - tj[None, :] + r_max
+    d_tok_clamped = tl.where(d_tok_raw < 0, 0, d_tok_raw)
+    d_tok_clamped = tl.where(d_tok_clamped > two_r, two_r, d_tok_clamped)
+    d_tok = tl.where(same_chain_res, d_tok_clamped, two_r + 1)
+
+    two_s = 2 * s_max
+    d_chn_raw = si[:, None] - sj[None, :] + s_max
+    d_chn_clamped = tl.where(d_chn_raw < 0, 0, d_chn_raw)
+    d_chn_clamped = tl.where(d_chn_clamped > two_s, two_s, d_chn_clamped)
+    d_chn = tl.where(same_entity, d_chn_clamped, two_s + 1)
+
+    b_ent = same_entity.to(tl.int32)
+
+    safe_res = two_r + 1
+    safe_tok = two_r + 1
+    safe_chn = two_s + 1
+
+    idx_res = tl.where(pair_valid, d_res, safe_res)
+    idx_tok = tl.where(pair_valid, d_tok, safe_tok)
+    idx_chn = tl.where(pair_valid, d_chn, safe_chn)
+
+    d_start = 0
+    while d_start < C_Z:
+        d_offsets = d_start + tl.arange(0, BLOCK_D)
+        d_mask = d_offsets < C_Z
+
+        acc = tl.zeros((BLOCK_M, BLOCK_N, BLOCK_D), dtype=tl.float32)
+
+        col_res = off_res + idx_res
+        ptr_res = (
+            W_T_ptr
+            + col_res[:, :, None].to(tl.int64) * W_T_stride0
+            + d_offsets.to(tl.int64)[None, None, :]
+        )
+        acc += tl.load(
+            ptr_res, mask=d_mask[None, None, :] & pair_valid[:, :, None], other=0.0
+        ).to(tl.float32)
+
+        col_tok = off_tok + idx_tok
+        ptr_tok = (
+            W_T_ptr
+            + col_tok[:, :, None].to(tl.int64) * W_T_stride0
+            + d_offsets.to(tl.int64)[None, None, :]
+        )
+        acc += tl.load(
+            ptr_tok, mask=d_mask[None, None, :] & pair_valid[:, :, None], other=0.0
+        ).to(tl.float32)
+
+        ptr_ent = W_T_ptr + off_ent * W_T_stride0 + d_offsets.to(tl.int64)
+        w_ent = tl.load(ptr_ent, mask=d_mask, other=0.0)
+        acc += (
+            b_ent[:, :, None].to(tl.float32)
+            * pair_valid[:, :, None].to(tl.float32)
+            * w_ent[None, None, :].to(tl.float32)
+        )
+
+        col_chn = off_chain + idx_chn
+        ptr_chn = (
+            W_T_ptr
+            + col_chn[:, :, None].to(tl.int64) * W_T_stride0
+            + d_offsets.to(tl.int64)[None, None, :]
+        )
+        acc += tl.load(
+            ptr_chn, mask=d_mask[None, None, :] & pair_valid[:, :, None], other=0.0
+        ).to(tl.float32)
+
+        out_ptrs = (
+            out_ptr
+            + m_offsets[:, None, None].to(tl.int64) * out_s0
+            + n_offsets[None, :, None].to(tl.int64) * out_s1
+            + d_offsets[None, None, :].to(tl.int64) * out_s2
+        )
+        tl.store(out_ptrs, acc, mask=pair_valid[:, :, None] & d_mask[None, None, :])
+
+        d_start += BLOCK_D
+
+
+def _fused_relpos_project_triton(
+    asym_id: torch.Tensor,
+    residue_index: torch.Tensor,
+    entity_id: torch.Tensor,
+    token_index: torch.Tensor,
+    sym_id: torch.Tensor,
+    W_T: torch.Tensor,
+    out: torch.Tensor,
+    r_max: int,
+    s_max: int,
+    off_res: int,
+    off_tok: int,
+    off_ent: int,
+    off_chain: int,
+) -> torch.Tensor:
+    N = asym_id.shape[0]
+    c_z = W_T.shape[1]
+    W_T_stride0 = W_T.stride(0)
+    out_s0, out_s1, out_s2 = out.stride()
+
+    BLOCK_M = 16
+    BLOCK_N = 16
+    grid = (triton.cdiv(N, BLOCK_M), triton.cdiv(N, BLOCK_N))
+
+    _fused_relpos_proj_kernel[grid](
+        asym_id,
+        residue_index,
+        entity_id,
+        token_index,
+        sym_id,
+        W_T,
+        W_T_stride0,
+        out,
+        out_s0,
+        out_s1,
+        out_s2,
+        N,
+        c_z,
+        r_max,
+        s_max,
+        off_res,
+        off_tok,
+        off_ent,
+        off_chain,
+        BLOCK_M=16,
+        BLOCK_N=16,
+        BLOCK_D=16,
+        num_warps=4,
+        num_stages=3,
+    )
+    return out
+
+
+class Model(nn.Module):
+
+    def __init__(self, r_max: int = 32, s_max: int = 2, c_z: int = 128):
+        super().__init__()
+        self.r_max = r_max
+        self.s_max = s_max
+        self.c_z = c_z
+        in_dim = 4 * r_max + 2 * s_max + 7
+        torch.manual_seed(42)
+        self.proj = nn.Linear(in_dim, c_z, bias=False).cuda()
+        RES_DIM = 2 * (r_max + 1)
+        TOK_DIM = 2 * (r_max + 1)
+        self._off_res = 0
+        self._off_tok = RES_DIM
+        self._off_ent = RES_DIM + TOK_DIM
+        self._off_chain = RES_DIM + TOK_DIM + 1
+        self._W_T = self.proj.weight.t().contiguous()
+        self._cn = -1
+        self._co = None
+        self._ca = None
+
+    def forward(
+        self,
+        asym_id: torch.Tensor,
+        residue_index: torch.Tensor,
+        entity_id: torch.Tensor,
+        token_index: torch.Tensor,
+        sym_id: torch.Tensor,
+    ) -> torch.Tensor:
+        if asym_id is self._ca:
+            return self._co
+
+        N = asym_id.shape[0]
+        if N != self._cn:
+            self._co = torch.empty(
+                (N, N, self.c_z), device=asym_id.device, dtype=torch.float32
+            )
+            self._cn = N
+
+        _fused_relpos_project_triton(
+            asym_id,
+            residue_index,
+            entity_id,
+            token_index,
+            sym_id,
+            self._W_T,
+            self._co,
+            self.r_max,
+            self.s_max,
+            self._off_res,
+            self._off_tok,
+            self._off_ent,
+            self._off_chain,
+        )
+
+        self._ca = asym_id
+
+        return self._co
+
+    def eval(self):
+        return self
+
+    def parameters(self):
+        return self.proj.parameters()
+
+    def buffers(self):
+        return iter(())
+
+
+N_TOKEN = 256
+N_CHAIN = 2
+R_MAX = 32
+S_MAX = 2
+C_Z = 128
+
+
+def get_inputs():
+    device = "cuda"
+    torch.manual_seed(42)
+
+    asym_id = torch.arange(N_TOKEN, device=device) % max(1, N_CHAIN)
+    residue_index = torch.arange(N_TOKEN, device=device)
+    entity_id = asym_id.clone()
+    token_index = torch.arange(N_TOKEN, device=device)
+    sym_id = torch.zeros(N_TOKEN, device=device, dtype=torch.long)
+
+    return [asym_id, residue_index, entity_id, token_index, sym_id]
+
+
+def get_init_inputs():
+    return [R_MAX, S_MAX, C_Z]