Initialize legacy kernels compatibility mirror
Browse files
README.md
CHANGED
|
@@ -1,116 +1,9 @@
|
|
| 1 |
-
# fp4-gemm
|
| 2 |
|
| 3 |
-
|
|
|
|
| 4 |
|
| 5 |
-
|
| 6 |
-
buffers and produces BF16 output. It is designed to pair with
|
| 7 |
-
`flashrt/fp4-fused-ops` and other static low-bit transformer/diffuser runtime
|
| 8 |
-
paths.
|
| 9 |
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
- `sfa_size_bytes(rows, dim)`
|
| 13 |
-
- `quantize_fp4_sfa_fp16(x, packed=None, sfa=None, is_sfb=False)`
|
| 14 |
-
- `quantize_fp4_sfa_bf16(x, packed=None, sfa=None, is_sfb=False)`
|
| 15 |
-
- `dequantize_fp4_sfa_fp16(packed, sfa, out=None, is_sfb=False)`
|
| 16 |
-
- `nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0, out=None, variant=-1)`
|
| 17 |
-
- `nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out=None)`
|
| 18 |
-
- `nvfp4_gemm_bias_residual_bf16(a_packed, b_packed, sfa, sfb, bias, residual, out=None)`
|
| 19 |
-
- `nvfp4_gemm_residual_bf16(a_packed, b_packed, sfa, sfb, residual, alpha=1.0, out=None)`
|
| 20 |
-
- `nvfp4_gemm_bias_gelu_bf16(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out=None)`
|
| 21 |
-
- `nvfp4_gemm_bias_gelu_nvfp4(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out_packed=None, out_sfa=None)`
|
| 22 |
-
- `nvfp4_gemm_streamk_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0, out=None)`
|
| 23 |
-
- `nvfp4_gemm_streamk_bias_bf16(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out=None)`
|
| 24 |
-
- `fp4_w4a16_linear_bf16(...)` is retained as a compatibility alias
|
| 25 |
-
|
| 26 |
-
## Tensor Contract
|
| 27 |
-
|
| 28 |
-
- `a_packed`: `torch.uint8`, shape `(M, K / 2)`.
|
| 29 |
-
- `b_packed`: `torch.uint8`, shape `(N, K / 2)`.
|
| 30 |
-
- `sfa`: `torch.uint8`, CUTLASS SFA layout for `(M, K)`.
|
| 31 |
-
- `sfb`: `torch.uint8`, CUTLASS SFB layout for `(N, K)`.
|
| 32 |
-
- output: `torch.bfloat16`, shape `(M, N)`.
|
| 33 |
-
- `K` must be divisible by 16.
|
| 34 |
-
- Targets: Blackwell `sm_110a` (Jetson AGX Thor, CUDA 13+) and `sm_120a`
|
| 35 |
-
(RTX Blackwell, CUDA 12.8+).
|
| 36 |
-
|
| 37 |
-
`variant` selects the CUTLASS schedule:
|
| 38 |
-
|
| 39 |
-
- `-1`: architecture-aware auto-dispatch (public default).
|
| 40 |
-
- `0`: default `<128,128,256>` cooperative schedule.
|
| 41 |
-
- `1`: widen `<128,256,128>` schedule, intended for very large `N`.
|
| 42 |
-
- `2`: pingpong schedule for A/B testing shape-specific wins.
|
| 43 |
-
|
| 44 |
-
The canonical linear API and FP4/SFA quantize/dequantize helpers are available
|
| 45 |
-
on both SM110 and SM120. SM110 additionally provides the GROOT N1.7 production
|
| 46 |
-
epilogues `nvfp4_gemm_bias_bf16`, `nvfp4_gemm_bias_residual_bf16`, and
|
| 47 |
-
`nvfp4_gemm_bias_gelu_nvfp4`. The latter emits packed FP4 plus CUTLASS SFA so
|
| 48 |
-
the following projection can consume it without a BF16 materialization and a
|
| 49 |
-
standalone quantization launch. Stream-K and the older BF16 GELU epilogue keep
|
| 50 |
-
their existing SM120 dispatch and reject unsupported architectures explicitly.
|
| 51 |
-
|
| 52 |
-
The SM110 release gate includes the production `(M,N,K)` shapes
|
| 53 |
-
`(41,4608,1536)`, `(41,6144,1536)`, and `(41,1536,6144)`, plus the legacy
|
| 54 |
-
`M=51` compatibility row. The kernels are the native sources used by FlashRT's
|
| 55 |
-
GROOT N1.7 Thor NVFP4 pipeline.
|
| 56 |
-
|
| 57 |
-
## Minimal Usage
|
| 58 |
-
|
| 59 |
-
```python
|
| 60 |
-
from kernels import get_kernel
|
| 61 |
-
import torch
|
| 62 |
-
|
| 63 |
-
ops = get_kernel("flashrt/fp4-gemm", version=1, trust_remote_code=True)
|
| 64 |
-
|
| 65 |
-
x = torch.randn((32, 256), device="cuda", dtype=torch.float16)
|
| 66 |
-
w = torch.randn((512, 256), device="cuda", dtype=torch.float16)
|
| 67 |
-
|
| 68 |
-
a_packed, sfa = ops.quantize_fp4_sfa_fp16(x, is_sfb=False)
|
| 69 |
-
b_packed, sfb = ops.quantize_fp4_sfa_fp16(w, is_sfb=True)
|
| 70 |
-
|
| 71 |
-
y = ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0)
|
| 72 |
-
```
|
| 73 |
-
|
| 74 |
-
For BF16 model activations, use the direct producer so the hot path does not
|
| 75 |
-
materialize an intermediate FP16 tensor:
|
| 76 |
-
|
| 77 |
-
```python
|
| 78 |
-
x_bf16 = torch.randn((1, 5120), device="cuda", dtype=torch.bfloat16)
|
| 79 |
-
a_packed, sfa = ops.quantize_fp4_sfa_bf16(x_bf16)
|
| 80 |
-
```
|
| 81 |
-
|
| 82 |
-
The BF16 entry writes the same E2M1 bytes and CUTLASS SFA/SFB layout as
|
| 83 |
-
`quantize_fp4_sfa_fp16(x_bf16.to(torch.float16))` for finite FP16-range
|
| 84 |
-
inputs. It is an additive API; the existing FP16 producer remains unchanged.
|
| 85 |
-
|
| 86 |
-
The quantize/dequantize helpers are included for examples and validation. A
|
| 87 |
-
production runtime should keep weights prepacked and should avoid quantizing in
|
| 88 |
-
the hot path unless that producer kernel is part of the intended low-bit block.
|
| 89 |
-
|
| 90 |
-
Use the bias/GELU and residual variants to avoid returning to BF16
|
| 91 |
-
elementwise code between low-bit GEMMs. Stream-K variants are selected only
|
| 92 |
-
for the validated large down-projection shapes; unsupported shapes reject
|
| 93 |
-
rather than silently selecting a losing schedule.
|
| 94 |
-
|
| 95 |
-
## Validation
|
| 96 |
-
|
| 97 |
-
```bash
|
| 98 |
-
python fp4-gemm/tests/test_fp4_gemm.py --backend source --mode full
|
| 99 |
-
python fp4-gemm/tests/test_fp4_gemm.py --backend installed --mode full \
|
| 100 |
-
--artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux
|
| 101 |
-
python fp4-gemm/benchmarks/benchmark.py --backend installed --mode headline \
|
| 102 |
-
--artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux
|
| 103 |
-
|
| 104 |
-
# Thor model-shape gate
|
| 105 |
-
python fp4-gemm/tests/test_fp4_gemm.py --backend installed \
|
| 106 |
-
--mode thor-models \
|
| 107 |
-
--artifact fp4-gemm/build/torch211-cxx11-cu130-aarch64-linux
|
| 108 |
-
```
|
| 109 |
-
|
| 110 |
-
The correctness reference dequantizes the same FP4/SFA and FP4/SFB inputs used
|
| 111 |
-
by the kernel, then computes the PyTorch GEMM reference from those dequantized
|
| 112 |
-
low-bit values.
|
| 113 |
-
|
| 114 |
-
The producer gate also checks the BF16 direct entry byte-for-byte against the
|
| 115 |
-
established FP16 compatibility chain at decode widths 5120, 6144 and 17408,
|
| 116 |
-
plus multi-row activation and SFB layouts.
|
|
|
|
| 1 |
+
# flashrt/fp4-gemm
|
| 2 |
|
| 3 |
+
This repository is a compatibility mirror for older `kernels` clients
|
| 4 |
+
that resolve repositories through the default Hugging Face model repo API.
|
| 5 |
|
| 6 |
+
Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/fp4-gemm
|
|
|
|
|
|
|
|
|
|
| 7 |
|
| 8 |
+
Do not edit this mirror by hand. It is generated from the Kernel Hub
|
| 9 |
+
`vN` branches and contains the same `build/**` artifacts.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|