import torch
import torch.nn.functional as F
from triton.testing import do_bench
from torch._inductor import config as inductor_config
import math

def position_bias_softmax(scores, weight=None, pw_bias=False):
    scores = scores.to(torch.float32)
    context_position = torch.arange(2048, dtype=torch.long, device="cuda")[:, None]
    memory_position = torch.arange(2048, dtype=torch.long, device="cuda")[None, :]
    relative_position = memory_position - context_position  # shape (query_length, key_length)
    relative_buckets = 0
    num_buckets=32
    max_distance=128
    relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
    max_exact = num_buckets // 2
    is_small = relative_position < max_exact
    relative_position_if_large = max_exact + (
        torch.log(relative_position.float() / max_exact)
        / math.log(max_distance / max_exact)
        * (num_buckets - max_exact)
    ).to(torch.long)
    relative_position_if_large = torch.min(
        relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)
    )

    relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
    values = F.embedding(relative_buckets, weight)
    values = values.permute([2, 0, 1]).unsqueeze(0) 
    scores = scores + values

    return F.softmax(scores, dim=-1).to(torch.float16)


scores = torch.randn(8, 2048, 2048, device="cuda", dtype=torch.float16)
weight = torch.randn(32, 1, device="cuda")
position_bias_softmax(scores, weight)
compiled = torch.compile(position_bias_softmax)

compiled(scores, weight=weight)
gb = 2 * scores.element_size() * scores.numel() / 1e9
sec = do_bench(lambda: compiled(scores, weight=weight)) / 1e3
print(f"weighted bias gb/s: {gb/sec}")