Skip to content

mlops-attention 0.2.0

Choose a tag to compare

@srose69 srose69 released this 23 Oct 00:14
· 13 commits to main since this release

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_prefill append a chunk of tokens to the cache and, with
    queries, return their causal attention over the whole cache (position tokens + i reads 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; encode and append run 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 in q8_0),
    prefill 2.1× faster at 32K, where llama.cpp fits only a q4_0 cache. 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.py runs llama-bench's tests on the latter.
  • C API: mlops_fbla_prefill; mlops_fbla_cache's pointers are no longer const (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).