kimbochen

#include <iostream>
#include <cuda_runtime.h>

int main() {
    // Ref: https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html
    // Compilation command: nvcc device_query.cpp -arch=sm_100 -o device_query && ./device_query
    int device = 0;
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, device);
    

kimbochen

import torch
import torch.nn.functional as F


@torch.compile()
def baseline(q_BHTD, k_BJTD, v_BJTD, gq_ratio):
    k_BHTD = k_BJTD.repeat_interleave(gq_ratio, 1)
    v_BHTD = v_BJTD.repeat_interleave(gq_ratio, 1)
    o_BHTD = F.scaled_dot_product_attention(q_BHTD, k_BHTD, v_BHTD, is_causal=True)
    return o_BHTD

kimbochen

import functools
from dataclasses import asdict, dataclass
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass

kimbochen

kimbochen

kimbochen

ML Efficiency Notes

GPU Specs

Name	FP16 Compute	Memory Bandwidth	Memory Size	TDP
A100	312 TFLOP/s	2 TB/s	40 GB	250 W
H100	750 TFLOP/s	2 TB/s	80 GB	350 W
A10	125 TFLOP/s	0.6 TB/s	24 GB	150 W

kimbochen

kimbochen

kimbochen

kimbochen