ChinmayK0607

import torch
import torch.nn as nn
from torch import einsum
import einops
from dataclasses import dataclass

@dataclass
class KDACache:
    S: torch.Tensor
    past_len: torch.Tensor | None = None