Compute Shaders con WebGPU para Inferencia ML: Ejecución de Modelos en el Cliente

¿Por Qué No Usar ONNX Runtime Web Directamente?

ONNX Runtime Web es una herramienta sólida para producción. Sin embargo, es una abstracción sobre WebGPU (o WebGL, o WASM), y las abstracciones tienen costos. Al perfilar una sesión de ONNX Runtime Web corriendo un modelo DistilBERT, aproximadamente el 18% del tiempo total de inferencia se iba en overhead de dispatch de operadores, bookkeeping de alocación de memoria, y pases de optimización de grafo que ya habían sido aplicados al modelo.

Escribir compute shaders directamente te da control sobre:

• Las dimensiones exactas de workgroup y la geometría de dispatch
• El layout de memoria en storage y uniform buffers
• La sincronización entre pases
• Aritmética cuantizada sin dequantización en runtime

La contra es que se está escribiendo código de GPU. No hay almuerzo gratis.

Armando el Pipeline WebGPU

Primero, el boilerplate. El código de negociación de adapter/device es boilerplate estándar; las partes interesantes vienen a continuación.

const adapter = await navigator.gpu.requestAdapter({
 powerPreference: 'high-performance'
});

const device = await adapter.requestDevice({
 requiredFeatures: ['timestamp-query'],
 requiredLimits: {
 maxStorageBufferBindingSize: 256 * 1024 * 1024,
 maxBufferSize: 256 * 1024 * 1024,
 maxComputeWorkgroupSizeX: 256,
 maxComputeInvocationsPerWorkgroup: 256,
 }
});

Ese maxStorageBufferBindingSize es importante, el default es 128MB en muchas implementaciones, y un modelo cuantizado de 125M parámetros puede excederlo fácilmente cuando incluís los buffers de KV cache.

El Núcleo: Multiplicación de Matrices en WGSL

La columna vertebral de la inferencia de transformers es la multiplicación de matrices. A continuación, va un compute shader WGSL que hace matmul con tiles y memoria compartida, que es el approach mínimo viable para rendimiento aceptable:

@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> result: array<f32>;

struct Dimensions {
 M: u32,
 N: u32,
 K: u32,
}
@group(0) @binding(3) var<uniform> dims: Dimensions;

const TILE_SIZE: u32 = 16u;

var<workgroup> tileA: array<array<f32, 16>, 16>;
var<workgroup> tileB: array<array<f32, 16>, 16>;

@compute @workgroup_size(16, 16, 1)
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(local_invocation_id) lid: vec3<u32>,
) {
 let row = gid.x;
 let col = gid.y;
 var sum: f32 = 0.0;

 let numTiles = (dims.K + TILE_SIZE - 1u) / TILE_SIZE;

 for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
 let tiledCol = t * TILE_SIZE + lid.y;
 let tiledRow = t * TILE_SIZE + lid.x;

 if (row < dims.M && tiledCol < dims.K) {
 tileA[lid.x][lid.y] = a[row * dims.K + tiledCol];
 } else {
 tileA[lid.x][lid.y] = 0.0;
 }

 if (tiledRow < dims.K && col < dims.N) {
 tileB[lid.x][lid.y] = b[tiledRow * dims.N + col];
 } else {
 tileB[lid.x][lid.y] = 0.0;
 }

 workgroupBarrier();

 for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
 sum = sum + tileA[lid.x][k] * tileB[k][lid.y];
 }

 workgroupBarrier();
 }

 if (row < dims.M && col < dims.N) {
 result[row * dims.N + col] = sum;
 }
}

Este es el approach tileado estándar. El detalle clave que la mayoría de los tutoriales se pierden: esas dos llamadas a workgroupBarrier() son innegociables. La primera asegura que el tile esté completamente cargado antes de computar; la segunda asegura que la computación esté completa antes de reusar la memoria del tile. Saquí cualquiera de las dos y se va a tener resultados incorrectos intermitentes que son un infierno para debuggear.

Yendo a Cuantizado: Matmul INT4

Correr matmul en FP32 funciona, pero es un desperdicio para inferencia. La mayoría de los modelos de producción usan cuantización INT4 o INT8. A continuación, resulta interesante, WGSL no tiene tipos int4 nativos, así que empaquetás dos valores int4 en un solo u8 (bueno, u32, ya que el tipo más pequeño de WGSL es 32-bit).

fn unpack_int4(packed: u32, index: u32) -> f32 {
 let shift = (index % 8u) * 4u;
 let nibble = (packed >> shift) & 0xFu;
 // Sign-extend desde 4-bit
 let signed = i32(nibble) - select(0, 16, nibble >= 8u);
 return f32(signed);
}

@compute @workgroup_size(256, 1, 1)
fn quantized_matmul(
 @builtin(global_invocation_id) gid: vec3<u32>,
) {
 let row = gid.x / dims.N;
 let col = gid.x % dims.N;

 if (row >= dims.M) { return; }

 var sum: f32 = 0.0;
 let groupSize: u32 = 32u;

 for (var k: u32 = 0u; k < dims.K; k = k + 1u) {
 let group_idx = k / groupSize;
 let scale = scales[row * (dims.K / groupSize) + group_idx];
 let zero = zeros[row * (dims.K / groupSize) + group_idx];

 let packed_idx = (row * dims.K + k) / 8u;
 let w = unpack_int4(weights_packed[packed_idx], k);
 let dequant = (w - zero) * scale;

 sum = sum + input[col * dims.K + k] * dequant;
 }

 result[row * dims.N + col] = sum;
}

La función unpack_int4 hace el trabajo pesado. Guardamos 8 valores int4 por u32, desempaquetamos on the fly, y aplicamos dequantización por grupo. El tamaño de grupo de 32 coincide con lo que usan GPTQ y AWQ, así que se puede cargar modelos cuantizados directamente.

El Mecanismo de Atención

Self-attention es la otra pieza crítica para rendimiento. Lo que la mayoría de la gente hace mal: implementan atención como tres dispatches separados (QK^T, softmax, attnV). Pero se puede fusionar el softmax en el mismo pase que la multiplicación Q*K^T si usás el truco de softmax online:

@compute @workgroup_size(256, 1, 1)
fn fused_attention(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(local_invocation_id) lid: vec3<u32>,
) {
 let head = gid.x / seq_len;
 let query_pos = gid.x % seq_len;

 if (head >= num_heads) { return; }

 var max_score: f32 = -1e10;
 var sum_exp: f32 = 0.0;
 var acc: array<f32, 64>; // head_dim

 // Online softmax: pase único sobre K
 for (var k: u32 = 0u; k < seq_len; k = k + 1u) {
 var score: f32 = 0.0;
 for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
 score += q[head * seq_len * head_dim + query_pos * head_dim + d]
 * keys[head * seq_len * head_dim + k * head_dim + d];
 }
 score = score / sqrt(f32(head_dim));

 let new_max = max(max_score, score);
 let correction = exp(max_score - new_max);
 sum_exp = sum_exp * correction + exp(score - new_max);

 for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
 acc[d] = acc[d] * correction
 + exp(score - new_max)
 * values[head * seq_len * head_dim + k * head_dim + d];
 }
 max_score = new_max;
 }

 // Normalizar
 for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
 output[head * seq_len * head_dim + query_pos * head_dim + d]
 = acc[d] / sum_exp;
 }
}

El softmax online evita un segundo pase sobre los scores de atención para normalizar. En una secuencia de 512 tokens con 12 heads, esta fusión ahorró ~4ms por capa en una RTX 3060, que se acumula rápido a lo largo de 6 o 12 capas.

La Orquestación de Dispatches

Atando todo junto en JavaScript. A continuación, es donde coordinás la inferencia multi-pase:

class WebGPUInferenceEngine {
 constructor(device, model) {
 this.device = device;
 this.pipelines = {};
 this.buffers = {};
 this.model = model;
 }

 async buildPipelines() {
 this.pipelines.embed = await this.createComputePipeline(embeddingShader);
 this.pipelines.layerNorm = await this.createComputePipeline(layerNormShader);
 this.pipelines.attention = await this.createComputePipeline(fusedAttentionShader);
 this.pipelines.ffn = await this.createComputePipeline(quantizedMatmulShader);
 this.pipelines.softmax = await this.createComputePipeline(finalSoftmaxShader);
 }

 runInference(inputTokens) {
 const encoder = this.device.createCommandEncoder();

 // Embedding lookup
 this.dispatchEmbedding(encoder, inputTokens);

 // Capas del transformer
 for (let layer = 0; layer < this.model.numLayers; layer++) {
 this.dispatchLayerNorm(encoder, layer, 'pre_attn');
 this.dispatchAttention(encoder, layer);
 this.dispatchLayerNorm(encoder, layer, 'pre_ffn');
 this.dispatchFFN(encoder, layer);
 }

 // Norm final + logits
 this.dispatchLayerNorm(encoder, -1, 'final');
 this.dispatchLogits(encoder);

 this.device.queue.submit([encoder.finish()]);
 }
}

Todos los dispatches van a un solo command encoder. El driver de la GPU maneja el tracking de dependencias entre dispatches basándose en el uso de buffers, esta es un área donde WebGPU realmente brilla comparado con hacer esto en CUDA, donde tendrías que manejar streams manualmente.

Benchmarks: WebGPU vs WebGL vs Server

Los corrí en una MacBook Pro 2024 (M3 Pro) y un desktop con RTX 3060, usando un DistilBERT cuantizado (66M params, INT4):

Método	M3 Pro	RTX 3060	Notas
WebGPU Compute (nuestro)	19ms	23ms	Shaders WGSL custom
ONNX Runtime Web (WebGPU)	31ms	38ms	Mismo modelo, config default
ONNX Runtime Web (WebGL)	67ms	54ms	Backend WebGL2
Server (Lambda + ONNX)	45ms*	45ms*	*Incluye ~30ms red
Server (Lambda cold start)	2.8s	2.8s	Penalidad de warm-up

Los números del M3 Pro son notables. La implementación de WebGPU de Apple es genuinamente excelente, su backend de Metal para WebGPU está maduro y bien optimizado. Los números de la RTX 3060 son ligeramente más altos porque el backend Dawn de Chrome agrega una capa fina de traducción a Vulkan.

La victoria real no es la latencia cruda, es la eliminación de cold starts y variabilidad de red. Un P99 de 25ms en el navegador supera a un P99 de 3 segundos de una función serverless que puede necesitar cold-start.

Trampas de Layout de Memoria

Un obstáculo común que vale la pena destacar: las reglas de alineación de storage buffers en WGSL son estrictas y están mal documentadas. Cada miembro del struct debe estar alineado a un múltiplo de su tamaño, y el struct en sí debe estar alineado al mayor alineamiento de sus miembros. Para un struct mixto f32/u32:

struct LayerParams {
 weight_offset: u32, // offset 0, align 4
 bias_offset: u32, // offset 4, align 4
 scale: f32, // offset 8, align 4
 num_elements: u32, // offset 12, align 4
 // Total: 16 bytes, align 4, todo bien
}

// Pero si agregars un vec4:
struct LayerParamsV2 {
 weight_offset: u32, // offset 0
 // 12 bytes de padding!
 params: vec4<f32>, // offset 16, align 16
 num_elements: u32, // offset 32
 // 12 bytes de padding!
 // Total: 48 bytes
}

Ese padding es silencioso y brutal. Tus escrituras de buffer del lado JavaScript van a estar mal si no lo se tiene en cuenta. Utilizar anotaciones offset_of o directamente evitar tipos vec en structs de uniform.

Cuándo NO Usar Este Approach

Hay limitaciones importantes a considerar. Este approach no es adecuado cuando:

• Tu modelo tiene más de ~500M parámetros. La memoria GPU en contextos de browser es limitada, y la se comparte con el renderer.
• Necesitás resultados determinísticos entre dispositivos. El comportamiento de punto flotante varía entre vendors de GPU de maneras que importan para algunas aplicaciones.
• Tus usuarios están en gráficos integrados. Intel UHD 630 puede correr estos shaders, pero a 10x la latencia. Detectar features y tener un fallback.
• Necesitás mantener los pese es del modelo privados. Todo lo que se envía al browser es inspeccionable. Si tu modelo es propietario, esto no es para vos.

Qué Viene Después

Un área activa de exploración es la implementación WebGPU de flash attention para ventanas de contexto más largas. El desafío es que WGSL no te da el control fino de shared memory que te da CUDA, así que la estrategia de tiling tiene que ser diferente. Esto valdrá la pena examinar cuando haya benchmarks significativos disponibles.

El punto más amplio es este: la industria pasó una década asumiendo que la inferencia ML pertenece a los servidores. Para una categoría creciente de modelos y case es de uso, la GPU del cliente no solo es adecuada, es superior. Cero varianza de latencia, cero cold starts, cero costos de servidor. El tooling finalmente se está poniendo a la altura para hacer esto práctico.