forked from ccf-ai-infra/TileOPs-Metax
34 lines
991 B
Python
34 lines
991 B
Python
from typing import Tuple
|
|
|
|
import torch
|
|
|
|
from workloads.base import WorkloadBase
|
|
|
|
|
|
class GLADecodeTest(WorkloadBase):
|
|
|
|
def __init__(
|
|
self,
|
|
batch: int,
|
|
heads: int,
|
|
dim_k: int,
|
|
dim_v: int,
|
|
dtype: torch.dtype,
|
|
scale: float = -1.0,
|
|
) -> None:
|
|
self.batch = batch
|
|
self.heads = heads
|
|
self.dim_k = dim_k
|
|
self.dim_v = dim_v
|
|
self.dtype = dtype
|
|
self.scale = scale
|
|
|
|
def gen_inputs(self) -> Tuple[torch.Tensor, ...]:
|
|
B, H, DK, DV = self.batch, self.heads, self.dim_k, self.dim_v
|
|
q = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1
|
|
k = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1
|
|
v = torch.randn(B, H, DV, device="cuda", dtype=self.dtype) * 0.1
|
|
gk = -torch.rand(B, H, DK, device="cuda", dtype=self.dtype)
|
|
state = torch.randn(B, H, DK, DV, device="cuda", dtype=self.dtype) * 0.1
|
|
return q, k, v, gk, state
|