diff --git a/dlblas/kernels/ks_competition/triton/QY2026_Autumn_Foxxiaohuli_Task04RelativePositionEncoding_MetaX.py b/dlblas/kernels/ks_competition/triton/QY2026_Autumn_Foxxiaohuli_Task04RelativePositionEncoding_MetaX.py new file mode 100644 index 000000000..1b0985b27 --- /dev/null +++ b/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]