youkaichao

import torch
from torch.utils.cpp_extension import load_inline

src = {
    "cuda": r"""
#include <cuda_runtime.h>
#include <torch/all.h>
#include <c10/cuda/CUDAStream.h>

__global__ void computation_kernel(unsigned long long total_nanosec) {

youkaichao

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

// Define the kernel with illegal memory access
__global__ void illegalWildPointerKernel(int* data, int size) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    __nanosleep(1000000000ULL);  // Sleep for 1 second
    int* wild_pointer = (int*)0x100;
    if (idx == 0) {

youkaichao

import torch
import torch.distributed as dist

use_nccl = False

dist.init_process_group(backend="nccl" if use_nccl else "gloo")

rank = dist.get_rank()
torch.cuda.set_device(rank % 8)

youkaichao

import torch.distributed as dist
import torch
import time

dist.init_process_group(backend="nccl")

rank = dist.get_rank()
torch.cuda.set_device(rank)

N_warmup = 10

youkaichao

import torch
import torch.nn as nn

# an Embedding module containing 10 tensors of size 3
embedding = nn.Embedding(10, 3)
embedding.weight.requires_grad_(False)
# a batch of 4 indices
input = torch.LongTensor([1, 2, 4, 5])
output = embedding(input)

youkaichao

import os
from typing import List
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch
import torch.distributed as dist

dist.init_process_group(backend="gloo")
rank = local_rank = dist.get_rank()
world_size = dist.get_world_size()
torch.cuda.set_device(local_rank)

youkaichao

import os
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn

def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                    out: torch.Tensor) -> None:
    out.copy_(q)

youkaichao

import os
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn

def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    out = q.clone()
    out += k

youkaichao

import os
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn

def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                    out: torch.Tensor) -> None:
    out.copy_(q)

youkaichao

import openai
import asyncio

async def get_choice_completion(prompt, choices):
    # Initialize an asynchronous OpenAI client
    async with openai.AsyncClient(base_url="http://127.0.0.1:8000/v1", api_key="abc") as client:
        choice_probs = {}

        # Calculate logprobs for each prompt + choice sequence
        for choice in choices: