Loom KernelsRust × CUDAGitHub

Documentation / Operators

A deliberately narrow surface.

An operator becomes supported only after contract, oracle, accelerator, named-baseline, and framework boundaries exist. Engine and end-to-end gates remain explicit even after the kernel is fast.

Supported

OperatorDtypesFusion boundary
RMSNormF32 · FP16 · BF16Standalone normalization
Add + RMSNormF32 · FP16 · BF16Residual update + normalization
RMSNorm + dynamic FP8F32 · FP16 · BF16 → E4M3FNNormalization + GEMM input quantization
SiLU-and-MulF32 · FP16 · BF16Split-half SwiGLU activation
SiLU-and-Mul + block FP8FP16 · BF16 → E4M3FNActivation + group-64/128 quantization
RoPE + paged-KV writeF32 · FP16 · BF16Packed Q/K rotation + native cache write
Greedy + sampled logprobF32 · FP16 · BF16Argmax + logsumexp + gather + tie rank
Selected-token logprob + rankF32 · FP16 · BF16Engine-owned sampling + one-pass normalization
Min-P filteringF32 · FP16 · BF16In-place row-max threshold; shape-gated in vLLM
Paged MQA/GQA decodeF32 · FP16 · BF16GQA packing + local split-K/LSE; short shapes route into vLLM
Greedy speculative verifyint32 drafts · int64 target IDsRagged acceptance + mismatch/bonus-token compaction

Next families

FP8 KV-cache compression

Reduce cache bytes and improve admitted context or batch size with quality and TPOT reported together.

Complete sampling tail

Own penalties, top-k/top-p, deterministic RNG, and top-k logprobs without host round trips.

KV-cache movement

Measure scheduler-facing block copy, gather, scatter, and compact work for prefix reuse and preemption.

Speculative extensions

Add tree, stochastic, or KV metadata only when a named profile exposes material non-GEMM cost.

Quantization plumbing

Remove scale, pack/unpack, dequant/requant, and layout passes around an unchanged vendor GEMM.

MoE routing + movement

Own routing, histogram/prefix sum, permutation, and combine while grouped GEMM stays vendor-owned.

Rust decode step

Prove zero-copy engine-neutral orchestration over borrowed tensors and streams without building an inference engine.

Implementation policy

  • Handwrite memory-bound and fusion-sensitive kernels.
  • Never implement dense, quantized, sparse, or grouped GEMM; keep the matrix core engine/vendor-owned.
  • Key tuning by device, dtype, shape, layout, and alignment.
  • Expose one stable Rust contract across backend implementations.
  • Do not turn incidental layout helpers into public APIs without a consumer.

Three admission questions

  1. Is the cost memory traffic, launch overhead, layout conversion, or scheduling metadata?
  2. Does a named inference-engine path leave a real boundary Loom can enter without copies?
  3. Can a real model or serving workload prove latency, memory, throughput, or goodput value?

Six independent gates

  1. Validated contract and invalid-input tests.
  2. Deterministic CPU or high-precision oracle.
  3. Accelerator correctness over edge and representative shapes.
  4. Warmed repeated measurement against a named baseline.
  5. Invocation from a real inference-engine execution path.
  6. TTFT, TPOT, throughput, memory, or goodput benefit.