msaroufim

Benchmarking

Offline

Each LLM can override:

• Max sequence length
• Extend tokens: amount of fresh compute; in decode, each step is usually 1 extend token
• CUDA graph max batch size: static buffers are captured per batch size from 1 to max, so larger values increase startup time and memory usage

Online benchmarking

msaroufim

"""
Benchmark: hl.dot_scaled vs nvfp4_matmul vs torch._scaled_mm vs torch.matmul
=============================================================================
Compares:
  1. torch.matmul (fp16) — cuBLAS baseline
  2. torch._scaled_mm (e4m3) — PyTorch native FP8 scaled matmul
  3. hl.dot_scaled (fp16) — hardware tl.dot_scaled, fp16 format
  4. hl.dot_scaled (e4m3) — hardware tl.dot_scaled, FP8 format
  5. nvfp4_matmul (sw dequant) — existing helion example, software FP4

msaroufim

Local LLM Service Setup

1. Install Ollama

curl -fsSL https://ollama.com/install.sh | sh

2. Download Model

ollama pull qwen3:32b

msaroufim

// Minimal NVRTC CUBIN generation example
// Compile: g++ -std=c++11 minimal_nvrtc_cubin.cpp -lnvrtc -lcuda -lcudart
// String -> cubin via nvrtc -> handle via the cuda driver API

#include <iostream>
#include <vector>
#include <nvrtc.h>
#include <cuda.h>

const char* kernelSource = R"(

msaroufim

import torch

# CUDA kernel with inline PTX
kernel_source = """
__global__ void vector_add(const float* a, const float* b, float* c, int n) {
    int idx;
    asm("mov.u32 %0, %%ctaid.x;" : "=r"(idx));
    int tid;
    asm("mov.u32 %0, %%tid.x;" : "=r"(tid));
    int ntid;

msaroufim

pytorch_operator,base_name,overload,folder_name,is_mapped
aten._adaptive_avg_pool2d.default,_adaptive_avg_pool2d,default,_adaptive_avg_pool2d,True
aten._adaptive_avg_pool2d_backward.default,_adaptive_avg_pool2d_backward,default,_adaptive_avg_pool2d_backward,True
aten._cudnn_rnn.default,_cudnn_rnn,default,_cudnn_rnn,True
aten._log_softmax.default,_log_softmax,default,_log_softmax,True
aten._log_softmax_backward_data.default,_log_softmax_backward_data,default,_log_softmax_backward_data,True
aten._softmax.default,_softmax,default,_softmax,True
aten._softmax_backward_data.default,_softmax_backward_data,default,_softmax_backward_data,True
aten._sparse_coo_tensor_with_dims_and_tensors.default,_sparse_coo_tensor_with_dims_and_tensors,default,_sparse_coo_tensor_with_dims_and_tensors,True
aten._to_copy.default,_to_copy,default,_to_copy,True

msaroufim

diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
index 938cb7dd97a..d3ac1369e6a 100644
--- a/torch/testing/_internal/common_methods_invocations.py
+++ b/torch/testing/_internal/common_methods_invocations.py
@@ -7443,6 +7443,57 @@ def reference_inputs_clone_contiguous(op, device, dtype, requires_grad, **kwargs
     yield SampleInput(a, kwargs={'memory_format': torch.channels_last_3d})
 
 
+def sample_inputs_copy(op_info, device, dtype, requires_grad, **kwargs):
+    """Sample inputs for copy and copy_ operations.

msaroufim

msaroufim

import torch
from torch import nn
from torch.distributed.tensor.placement_types import Replicate, Shard
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.parallel import parallelize_module


def dist_print(*args, **kwargs):

msaroufim

import torch
from torch.utils.cpp_extension import _get_cuda_arch_flags

def test_fix():
    print("Testing CUDA arch flags fix...")
    
    user_arch_flags = ['-gencode=arch=compute_86,code=sm_86']
    result = _get_cuda_arch_flags(user_arch_flags)
    
    print(f"User provided: {user_arch_flags}")