b12x
b12x is an SM120/SM121 CuTe DSL kernel library for (primarily) NVFP4 LLM
inference — the Blackwell GeForce and workstation parts (RTX 50-series,
RTX PRO) and GB10.
It is intentionally narrow. This is not a generic CUDA kernel collection or a
full model-serving stack, and it does not target any other architecture,
including SM100. It is a focused set of high-performance kernels plus the
runtime glue needed to launch them cleanly from sglang/vllm.
Install
pip install b12x
You need Python 3.10+, torch >= 2.12, and an SM120/SM121 GPU. The CuTe DSL
compiler and its CUDA 13 libraries come in as wheel dependencies
(nvidia-cutlass-dsl >= 4.5.2), so there is no build step — kernels are
JIT-compiled on first use and cached.
What's in here
GEMM (b12x/gemm/) — a dense block-scaled GEMM (DenseGemmKernel,
exposed as b12x::dense_gemm_launch) covering NVFP4 and MXFP8 operands with
BF16/FP16/FP32 outputs, plus fused linear layers on top of it: MXFP8
(b12x::mxfp8_linear_fused), 128x128 block-FP8
(b12x::block_fp8_linear_mxfp8_fused), and the grouped WO-projection paths
used by MLA attention output.
Attention (b12x/attention/) — contiguous (fixed-shape and packed-varlen)
and paged attention forward kernels, with BF16/FP16 and FP8 E4M3 KV caches,
GQA, sliding window, and attention sinks. Sparse MLA decode/prefill lives in
mla/, and the NSA/MSA logits indexer plus its top-k and scheduling kernels
in indexer/. Compressed MLA and GLM MLA/NSA are distinct contracts and kept
separate on purpose. paged/graph_replay.py has the metadata staging kernels
that make decode replayable under CUDA graphs.
MoE (b12x/moe/) — fused FP4 TP MoE in three flavors: a direct
micro-kernel decode path, a unified dynamic path (persistent grid, dynamic M
tiles, nvfp4/w4a8_mx/w4a8_nvfp4 weights), and W4A16 (BF16 activations
with inline FP4 weight dequant — no activation-scale math). SiLU, ReLU2, and
SwiGLU-OAI activations throughout.
Everything else — BF16→NVFP4 TMA quantization (b12x/quantization/), mHC
residual/projection kernels (b12x/integration/residual*.py), and an
IPC-backed PCIe one-shot allreduce (b12x/distributed/). The
b12x/integration/ layer is the boundary serving stacks talk to: it owns
planning, scratch layout, and policy, so integrations only supply metadata and
capacity limits.
Using it
Kernels are registered as torch custom ops under the b12x:: namespace, so
after import b12x they are callable as torch.ops.b12x.* and compose with
torch.compile and CUDA graphs. Higher-level Python entry points (kernel
classes, planners) live next to each kernel.
Compilation happens lazily per shape/config and is cached. For serving, warm up the shapes you need, then freeze:
import b12x
# ... run warmup traffic covering every shape you will serve ...
b12x.freeze_kernel_resolution("serving")
After the freeze, any request that would trigger a new kernel compile raises
KernelResolutionFrozenError instead of stalling a live request (or worse,
compiling inside CUDA graph capture).
Set B12X_PRINT_COMPILE_PROGRESS=1 to log each compiler invocation with its
cache-key parameters and duration — useful for figuring out what warmup
actually covered. B12X_TIMING=1 enables per-kernel timing logs.
Where to look next
tests/is the executable spec — every kernel has API and numerical reference tests showing exact tensor layouts and call sequences.benchmarks/has tuned invocations per kernel family (andprobe_*scripts from tile-sweep experiments).docs/has design notes: the MoE execution model, the eager-plan-bind architecture, and an SM120 MLA postmortem.
Failing that, ask your friendly neighborhood AI agent — it does fine here.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file b12x-0.30.2.tar.gz.
File metadata
- Download URL: b12x-0.30.2.tar.gz
- Upload date:
- Size: 1.1 MB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/6.2.0 CPython/3.14.5
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
aace20cf78b0f64b32731d3dac3169dbc9b767b1a11225468931a53651e7b422
|
|
| MD5 |
077d87de3904451fdddca5fae8a7f744
|
|
| BLAKE2b-256 |
6fbd8837a012124eb29d4077792b0214e08aad30053fb40fba9bc172cf28d4e3
|
File details
Details for the file b12x-0.30.2-py3-none-any.whl.
File metadata
- Download URL: b12x-0.30.2-py3-none-any.whl
- Upload date:
- Size: 940.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/6.2.0 CPython/3.14.5
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
081f6391f9a375dea8810bebf6a9d0276cacecde229b28715924f411e04bb38b
|
|
| MD5 |
d1b38e2e63fe520a82f8e041b2774a0f
|
|
| BLAKE2b-256 |
5120b97a6b844870c08aab83902bb545fa7a9643f63592c0a2e61be98b28df2f
|
