liangsu9988 commited on
Commit
ec63d06
·
verified ·
1 Parent(s): cc23714

Initialize legacy kernels compatibility mirror

Browse files
Files changed (1) hide show
  1. README.md +6 -113
README.md CHANGED
@@ -1,116 +1,9 @@
1
- # fp4-gemm
2
 
3
- FlashRT native Blackwell NVFP4 A4W4 GEMM kernels.
 
4
 
5
- This package consumes packed FP4 E2M1 tensors plus CUTLASS Sm1xx SFA/SFB scale
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
- ## Available Functions
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.