liangfu

# Compiler settings
CC = gcc
CFLAGS = -Wall -Wextra -O2
LDFLAGS = -lm

# Project files
SRC = rasterizer.c
OBJ = $(SRC:.c=.o)
TARGET = rasterizer

liangfu

import torch
import os
import torch_xla.core.xla_model as xm

def write_to_kv_cache(
        key: torch.Tensor,
        value: torch.Tensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        slot_mapping: torch.Tensor,

liangfu

import os
import pytest
import torch
import torch_neuronx
import torch_xla.core.xla_model as xm

@pytest.mark.parametrize("batch_size,seq_len,q_size,kv_size,use_torch_compile,disable_functionalization", [
    (2, 128, 32, 32, False, True),
    (2, 128, 32, 32, True, True),
    (2, 128, 32, 32, False, False),

liangfu

import pytest
import torch
import torch_xla.core.xla_model as xm

@pytest.mark.parametrize("batch_size,seq_len,q_size,kv_size", [
    (2, 128, 32, 32),
    (4, 256, 64, 64),
])
def test_split_consistency(batch_size, seq_len, q_size, kv_size):
    # Get XLA device

liangfu

import os
import time
import torch
import torch_xla.core.xla_model as xm

N = 16

def main():
    # os.environ["XLA_USE_BF16"] = "1"
    os.environ["NEURON_RT_STOCHASTIC_ROUNDING_EN"] = "1"

liangfu

import torch
import os
import depyf
import torch_xla.core.xla_model as xm

os.environ["NEURON_CC_FLAGS"]= " --model-type=transformer -O1 "
os.environ["NEURON_COMPILE_CACHE_URL"] = os.path.join(os.getcwd(), "_compile_cache")

@torch.compiler.allow_in_graph
def write_to_kv_cache(

liangfu

import time
import torch
import torch_xla.core.xla_model as xm

N = 128
n_iters = 100

def main():
    device = xm.xla_device()
    src = torch.arange(1, 2*N+1).reshape((2, N)).to(device=device)

liangfu

import torch
import torch_xla.core.xla_model as xm

@torch.compile(backend="openxla")
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

liangfu

from typing import Any, Dict, List, Optional, Tuple, Type

import torch
import torch_xla.core.xla_model as xm
import torch_xla.experimental.custom_kernel  # Required to register custom ops.

class PallasAttentionBackend:
    @torch.compile(backend="openxla")
    @staticmethod
    def copy_blocks(

liangfu

import boto3, json

question = """
How are you today?
"""

def main():
    session = boto3.Session()
    bedrock = session.client(service_name='bedrock-runtime', region_name="us-west-2")