gau-nernst

import torch
from torch import Tensor


def varlen_attn(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    cum_seq_q: Tensor,
    cum_seq_k: Tensor,

gau-nernst

uv venv --python=3.12 --managed-python
source .venv/bin/activate

uv pip install --group dev
uv pip install ninja  # or sudo apt install ninja-build
USE_DISTRIBUTED=0 USE_MKLDNN=0 BUILD_TEST=0 USE_FBGEMM=0 USE_NNPACK=0 USE_QNNPACK=0 USE_XNNPACK=0 USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 uv pip install --no-build-isolation -v -e .

gau-nernst

# modified from https://github.com/gpu-mode/kernelbot/blob/main/docker/amd-docker.Dockerfile

FROM ghcr.io/actions/actions-runner:latest

ENV CXX=clang++

ARG GPU_ARCH="gfx950"
ENV GPU_ARCH_LIST=$GPU_ARCH

RUN sudo apt-get update -y \