Repository navigation
mlops-attention 0.2.0
mlops-attention 0.2.0
Part of mlops: CUDA operations for NVIDIA GPUs sm_61 … sm_90.
Install
Nothing here downloads anything. Install first:
- PyTorch (CUDA build);
Then, in the directory with the downloaded wheel (cpXY = your Python, torchA.B / cuNNN = your PyTorch and its CUDA):
pip install --no-index -f . mlops-attention--no-index takes mlops-attention from this directory and its dependencies from what is installed, never from PyPI.
Other combinations: ./setup.sh attention from the source archive.
C / C++: unpack mlops-attention-0.2.0-linux-x86_64-cuda12.tar.xz (tar -xJf) and link lib/libmlops_attention.a (see its README).
Every artifact carries SASS for sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86, sm_87, sm_89, sm_90 (CUDA 12.0).
Assets
mlops-attention-0.2.0-linux-x86_64-cuda12.tar.xz(1012K)mlops-attention-0.2.0-src.tar.gz(124K)mlops_attention-0.2.0+cu126torch2.7-cp312-cp312-linux_x86_64.whl(5.6M)
Changes
- Prefill:
prefill(q, k, v, cache)/mlops_fbla_prefillappend a chunk of tokens to the cache and, with
queries, return their causal attention over the whole cache (positiontokens + ireads tokens
0 .. tokens + i). The groups the chunk closes are quantized on the GPU, bit for bit the encoder's records;
one streaming kernel computes the attention. Prefill is compute-bound, so it reads every closed group in full
(decode with every block refined, the same output in both modes): producer warps turn
the coarse and fine records into swizzled tiles, consumer warps run QK and PV by DP4A with an online softmax.
No scratch. empty(batch, kv_heads, capacity, mode)makes a cache to prefill into;encodeandappendrun on the GPU
through the same path.- Speed (GTX 1080): the prefill attention runs at 8 equivalent FP32 TFLOPS (4+2, 64K context), 2× PyTorch's
FP32 SDPA on Qwen3-0.6B's own q / k / v. Qwen3-8B-NVFP4 end to end against llama.cpp on the same weights
(llama-bench's tests): prefill 1.9× and decode 1.5× faster at a 16K context (llama.cpp's KV inq8_0),
prefill 2.1× faster at 32K, where llama.cpp fits only aq4_0cache. Table:docs/llamacpp-compare.md. - Samples:
02_streaming.py(one cache through prefill and decode),03_transformers.py(FBLA inside a
transformers model),04_generate.py(a complete generator without transformers' model: Qwen3 / Qwen2 / Llama
from safetensors, BF16 or ModelOpt NVFP4, NVFP4EffMA linear layers, mlops.positional RoPE, the layers that do
not fit streamed from host memory).bench/llama_bench.pyrunsllama-bench's tests on the latter. - C API:
mlops_fbla_prefill;mlops_fbla_cache's pointers are no longerconst(prefill writes the cache).
The pkg-config file now links libm: with 0.1.0's, a C program failed to link (log2f). The select kernels
are compiled once instead of in both format objects.
No dependency besides PyTorch (04_generate.py also uses nvfp4effma, mlops-positional, safetensors and the
model's tokenizer through transformers).