Practical GPU engineering: Metal shaders, profiling and edge deployment

Practical GPU engineering spans cross-platform shader programming, hardware counter diagnostics, and on-device edge deployment. Below is a deep dive into all three workflows.

1. Custom Apple Metal (MSL) Tiled Matrix Multiplication Shader

On Apple Silicon (M1–M4 / A-series), compute kernels are written in Metal Shading Language (MSL) and leverage Threadgroup Memory (Apple’s equivalent of CUDA Shared Memory) to avoid slow global memory reads.

Below is a tiled Matrix Multiplication (C = A × B) shader in MSL using 16 × 16 threadgroup tiles.

#include <metal_stdlib>
using namespace metal;

constant uint TILE_SIZE = 16;

struct MatrixDims {
    uint M; // Rows of A and C
    uint N; // Cols of B and C
    uint K; // Cols of A and Rows of B
};

kernel void matmul_tiled(
    device const float* A        [[buffer(0)]],
    device const float* B        [[buffer(1)]],
    device float* C              [[buffer(2)]],
    constant MatrixDims& dims    [[buffer(3)]],
    uint2 gid                    [[thread_position_in_grid]],
    uint2 tid                    [[thread_position_in_threadgroup]]
) {
    // Fast on-chip Threadgroup (SRAM) tiles
    threadgroup float tileA[TILE_SIZE][TILE_SIZE];
    threadgroup float tileB[TILE_SIZE][TILE_SIZE];

    float acc = 0.0f;
    uint numTiles = (dims.K + TILE_SIZE - 1) / TILE_SIZE;

    for (uint t = 0; t < numTiles; ++t) {
        uint tiledCol = t * TILE_SIZE + tid.x;
        uint tiledRow = t * TILE_SIZE + tid.y;

        // Collaboratively load 16x16 tiles into Threadgroup memory
        tileA[tid.y][tid.x] = (gid.y < dims.M && tiledCol < dims.K)
                              ? A[gid.y * dims.K + tiledCol] : 0.0f;
        tileB[tid.y][tid.x] = (tiledRow < dims.K && gid.x < dims.N)
                              ? B[tiledRow * dims.N + gid.x] : 0.0f;

        // Synchronize all threads in the threadgroup before reading tiles
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Compute partial dot product entirely inside registers/SRAM
        for (uint k = 0; k < TILE_SIZE; ++k) {
            acc += tileA[tid.y][k] * tileB[k][tid.x];
        }

        // Wait before overwriting tiles in the next loop iteration
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (gid.y < dims.M && gid.x < dims.N) {
        C[gid.y * dims.N + gid.x] = acc;
    }
}

• Key Hardware Mapping: [[thread_position_in_grid]] maps directly to CUDA's blockIdx * blockDim + threadIdx, while threadgroup_barrier(mem_flags::mem_threadgroup) maps to CUDA's __syncthreads(). For production FP16 workloads, Apple Silicon also exposes simdgroup_float8x8 hardware intrinsics to tap directly into Apple's matrix-multiply units.

2. Step-by-Step ncu Roofline Optimisation Walkthrough

Suppose you write a naive CUDA or Triton kernel for Row-Wise LayerNorm on a (4096, 4096) FP16 tensor, and it runs noticeably slower than PyTorch's native implementation.

Step 1: Capture Hardware Counters

Run NVIDIA Nsight Compute (ncu) targeting only your custom kernel name to avoid profiling overhead on setup code:

ncu --kernel-name regex:layernorm_naive --metrics \
sm__throughput.avg.pct_of_peak_sustained_elapsed,\
dram__throughput.avg.pct_of_peak_sustained_elapsed,\
l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\
smsp__warp_issue_stalled_long_scoreboard_per_warp_active.pct \
python run_layernorm.py

Step 2: Diagnose the Bottleneck on the Roofline

Inspect the output metrics:

• Compute (SM) Throughput: 11.4% of peak.

• Memory (DRAM/HBM) Throughput: 84.2% of peak.

• Warp Stall Reason (long_scoreboard): 68.5% (warps are stalled waiting for global HBM memory reads to return).

• Diagnosis: The kernel sits squarely under the sloped Memory-Bandwidth Roof on the Roofline chart. Furthermore, checking l1tex__t_sectors reveals uncoalesced memory reads—threads in the same warp are striding down columns instead of reading contiguous row addresses, wasting 75% of every 128-byte HBM cache line fetched.

Step 3: Apply Targeted Hardware Fixes

1. Fix Memory Coalescing: Ensure adjacent threads (threadIdx.x) read adjacent elements along the contiguous inner dimension (n_cols) so 32 FP16 reads merge into a single 64-byte coalesced transaction.

2. Vectorised Loads (float4 / tl.load block): Load 128 bits (8 FP16 values) per instruction instead of 16 bits, reducing total memory instruction count by 8×.

3. Warp-Shuffle Reductions: Replace slow Shared Memory reductions with register-to-register warp shuffles (__shfl_xor_sync in CUDA or tl.sum in Triton) to compute the row mean and variance without synchronization barriers.

Step 4: Verify on the Roofline

Re-running ncu confirms:

• Global Load Sectors: Drops by 4× (100% cache-line utilization).

• Long Scoreboard Stalls: Drops from 68.5% to < 15%.

• Wall-Clock Latency: Improves by 3.8×, pushing effective bandwidth above 85% of the H100's physical HBM limit.

3. Edge Deployment Pipeline: Quantizing & Running an LLM on Mobile/Laptop

Deploying an 8B-parameter LLM onto a laptop or smartphone (8 GB–16 GB shared RAM) requires compressing weights to 4-bit integers and executing through a zero-overhead C++ runtime like llama.cpp or Apple MLX.

[HF Safetensors (FP16, 16 GB)] 
       │  1. Convert & Apply Calibration Data (Imatrix)
       ▼
[GGUF Quantized Binary (Q4_K_M, 4.9 GB)]
       │  2. Memory-Map (mmap) into Unified RAM
       ▼
[Metal / Vulkan Compute Shaders (On-the-Fly Dequant + Matrix Mul)]

Step 1: Convert & Calibrate (Importance Matrix)

Raw 4-bit rounding degrades reasoning and coding accuracy. Using llama.cpp, first convert Hugging Face weights to an uncompressed GGUF file, then compute an Importance Matrix (imatrix) on a small calibration text sample to measure which weight channels have the largest activation impact:

python convert_hf_to_gguf.py ./Llama-3.1-8B --outfile llama-8b-f16.gguf
./llama-imatrix -m llama-8b-f16.gguf -f calibration.txt -o imatrix.dat

Step 2: Mixed-Bit K-Quantization (Q4_K_M)

Quantize the model using Q4_K_M (4-bit K-quant Medium), which keeps sensitive attention value/output layers at 6-bit (Q6_K) while compressing dense feed-forward layers to 4-bit (Q4_K) in super-blocks of 256 weights:

./llama-quantize --imatrix imatrix.dat llama-8b-f16.gguf llama-8b-q4_k_m.gguf Q4_K_M

This shrinks the model from 16.0 GB down to ~4.9 GB with less than 1% perplexity degradation.

Step 3: Zero-Copy Unified Memory Execution on Device

When running on an iPhone, Android device, or MacBook:

• Instant Load via mmap: The runtime memory-maps (mmap) the 4.9 GB .gguf file directly from flash storage into Unified Memory—eliminating RAM duplication and loading the model in under 1 second.

• Fused Dequantization Shaders: During token generation, custom Metal (ggml-metal.metal) or Vulkan (mul_mat_q4_K.comp) compute shaders load the packed 4-bit weight blocks directly into GPU registers, unpack them using bitwise shifts, and multiply them against FP16 activations on the fly—achieving 40–70+ tokens/sec on modern laptop GPUs and 15–25 tokens/sec on flagship mobile phones.