Hugging Face Kernels

The kernels package provides a standardized ecosystem for structuring, building, distributing, and loading compute kernels. It aims to solve several common developer pain points:

• Standardized Builds: Uses a builder component to take kernel source in a pre-defined layout with declarative configuration to produce compiled kernels for multiple backends (CUDA, ROCM, XPU), OSs, and architectures.
• Reproducibility: Enforces consistent build environments and steps.
• Efficient Distribution: Kernels are hosted on the Hugging Face Hub and loaded via the kernels Python package. This allows the package to automatically fetch the correct build for the user's specific system, avoiding lengthy local compilation times.
• Version Management: Supports loading multiple versions of the same kernel to prevent dependency conflicts and ensures compatibility across different PyTorch versions and hardware capabilities.
• Hub Integration: Leverages Hugging Face Hub features like XET for fast transfers, seamless versioning, and hardware compatibility visibility.

Every CPU kernel implementation should follow this directory and file organization to ensure proper compilation and Python integration:

my_kernel/
├── my_kernel_cpu/
│   ├── cpu_features.hpp          # CPUID detection (in its own namespace)
│   ├── my_kernel_cpu.cpp         # Dispatcher (handles AVX512 vs fallback)
│   ├── my_kernel_cpu.hpp         # Shared declarations
│   ├── my_kernel_cpu_torch.cpp   # Python ↔ C++ bridge
│   ├── my_kernel_avx512.cpp      # AVX512 implementation
│   └── my_kernel_avx512.hpp      # AVX512 declarations
├── torch-ext/
│   └── torch_binding.cpp         # PyTorch operator registration
└── build.toml                    # Multi-target compilation config

my_kernel/
├── my_kernel_cpu/
│   ├── cpu_features.hpp          # CPUID detection (own namespace)
│   ├── my_kernel_cpu.cpp         # Dispatcher
│   ├── my_kernel_cpu.hpp         # Shared declarations
│   ├── my_kernel_cpu_torch.cpp   # Python ↔ C++ bridge
│   ├── my_kernel_avx512.cpp      # AVX512 implementation
│   └── my_kernel_avx512.hpp      # AVX512 declarations
├── torch-ext/
│   └── torch_binding.cpp         # Op registration
└── build.toml                    # Multi-target compilation

To avoid slow double-compile/benchmark cycles on XPU, prune autotune configurations that are likely to cause register spills. Use an early_config_prune function to estimate per-thread GRF pressure using the following model:

regs ~= (accumulators * BM * BN + stages * (BM * BK + weight_tiles * BN * BK)) / (num_warps * warp_size)

Implementation Guidelines:

• Apply only to XPU backends.
• BM, BN, and BK should be read from autotune config kwargs or launch args.
• Accumulator count: Usually 2 for gate/up fused kernels, 1 for down/simple GEMM.
• Weight tile count: Usually 2 for gate/up fused kernels, 1 for down/simple GEMM.
• Stages: Use num_stages for pipelined MMA paths; use 1 for scalar/non-pipelined paths.
• Always keep a fallback configuration (e.g., the one with the lowest estimated register pressure).
• Use this as a coarse filter for catastrophic spillers, not a perfect predictor.

When writing optimized kernels for H100 (SM 9.0), use these patterns for performance:

Vectorized Memory Access

Use specialized types to load multiple elements in a single instruction:

• BFloat16: Use __nv_bfloat162 for 32-bit loads.
• FP16: Use __half2.
• FP32: Use float4.

Warp Shuffle Reductions

Use __shfl_xor_sync for efficient intra-warp reductions.

Thread Configuration

• Element-wise (RoPE, GEGLU): Use a standard BLOCK_SIZE (e.g., 256).
• Reductions (LayerNorm, RMSNorm): Ensure threads are rounded to a warp boundary and account for vectorized access (e.g., hidden_size / 2 for bf16/fp16).

// BFloat16 vectorization example
const __nv_bfloat162* vec_input = reinterpret_cast<const __nv_bfloat162*>(row_input);

#pragma unroll 4
for (int i = tid; i < vec_hidden; i += stride) {
    __nv_bfloat162 v = vec_input[i];
    float v0 = __bfloat162float(v.x);
    float v1 = __bfloat162float(v.y);
    sum_sq += v0 * v0 + v1 * v1;
}

// Warp shuffle reduction example
template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}