diff --git a/python/freetoken/kernel/triton/bf16_gemv.py b/python/freetoken/kernel/triton/bf16_gemv.py new file mode 100644 index 0000000000..2fc7fb107f --- /dev/null +++ b/python/freetoken/kernel/triton/bf16_gemv.py @@ -0,0 +1,58 @@ +"""One activation row times a bf16 weight: the GEMV a batch-1 decode step runs for every projection. + +The weight streams from HBM once, upcast to fp32 in registers and accumulated in fp32, so the result +is F.linear's up to the fp32 summation order. At one row the projection is bound by how fast the +weight streams in, and cuBLAS runs it on tensor-core GEMM tiles that leave part of the bandwidth +unused; this kernel reads nothing but the weight and the row. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _bf16_gemv_kernel( + x_ptr, w_ptr, out_ptr, N, K, + stride_wn, stride_wk, + BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, +): + pid = tl.program_id(0) + offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < N + w_row = w_ptr + offs_n[:, None] * stride_wn + acc = tl.zeros((BLOCK_N,), tl.float32) + for k0 in range(0, K, BLOCK_K): + offs_k = k0 + tl.arange(0, BLOCK_K) + k_mask = offs_k < K + w = tl.load(w_row + offs_k[None, :] * stride_wk, + mask=n_mask[:, None] & k_mask[None, :], other=0.0).to(tl.float32) + xk = tl.load(x_ptr + offs_k, mask=k_mask, other=0.0).to(tl.float32) + acc += tl.sum(w * xk[None, :], axis=1) + tl.store(out_ptr + offs_n, acc.to(out_ptr.dtype.element_ty), mask=n_mask) + + +def bf16_gemv(x: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """``x @ weight.T`` for a single row: ``x`` ``[..., K]`` holding one row, ``weight`` ``[N, K]`` + bf16; returns ``[..., N]`` in ``out_dtype``.""" + *lead, K = x.shape + N = weight.shape[0] + # F.linear raises on a mismatch; the kernel would read across rows, or past the weight's end + assert weight.shape[1] == K, f"x has K={K}, the weight is {tuple(weight.shape)}" + x1 = x.reshape(K).contiguous() + out = torch.empty(N, dtype=out_dtype, device=x.device) + # two rows a program spreads even a small GEMV over every SM, and a K tile up to 4096 wide + # keeps most rows to one pass + BLOCK_N = 2 + BLOCK_K = min(triton.next_power_of_2(K), 4096) + _bf16_gemv_kernel[(triton.cdiv(N, BLOCK_N),)]( + x1, weight, out, N, K, + weight.stride(0), weight.stride(1), + BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=4, + ) + return out.reshape(*lead, N) + + +__all__ = ["bf16_gemv"] diff --git a/python/freetoken/kernel/triton/dsv4/bf16_linear.py b/python/freetoken/kernel/triton/dsv4/bf16_linear.py index ee022addd1..51e89a9139 100644 --- a/python/freetoken/kernel/triton/dsv4/bf16_linear.py +++ b/python/freetoken/kernel/triton/dsv4/bf16_linear.py @@ -1,43 +1,19 @@ -"""bf16-weight GEMV with on-chip upcast + fp32 accumulate/output (DeepSeek-V4). +"""bf16-weight linear with fp32 math and fp32 output (DeepSeek-V4). Several DSV4 decode ops need fp32 *math* on bf16 weights (the compressor wkv/wgate gated pool, the MoE router) for numerical stability. Doing it as ``F.linear(x.float(), w.float())`` materializes an fp32 copy of the weight in HBM every step (read bf16 + write fp32 + re-read fp32) and runs a heavier fp32 GEMM. - -This kernel instead streams the bf16 weight from HBM once, upcasts to fp32 in -registers (on-chip), and accumulates in fp32 -> fp32 output. Same precision as the -reference (only the fp32 accumulation *order* differs, ~1e-6), at the HBM cost of a -plain bf16 read. Decode is M==1 (a GEMV); M>1 (prefill) falls back to F.linear. +Decode is M==1, a GEMV that upcasts the bf16 weight in registers instead; M>1 +(prefill) falls back to F.linear. """ from __future__ import annotations import torch import torch.nn.functional as F -import triton -import triton.language as tl - -@triton.jit -def _bf16_gemv_fp32_kernel( - x_ptr, w_ptr, out_ptr, N, K, - stride_wn, stride_wk, - BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, -): - pid = tl.program_id(0) - offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N) - n_mask = offs_n < N - w_row = w_ptr + offs_n[:, None] * stride_wn - acc = tl.zeros((BLOCK_N,), tl.float32) - for k0 in range(0, K, BLOCK_K): - offs_k = k0 + tl.arange(0, BLOCK_K) - k_mask = offs_k < K - w = tl.load(w_row + offs_k[None, :] * stride_wk, - mask=n_mask[:, None] & k_mask[None, :], other=0.0).to(tl.float32) - xk = tl.load(x_ptr + offs_k, mask=k_mask, other=0.0).to(tl.float32) - acc += tl.sum(w * xk[None, :], axis=1) - tl.store(out_ptr + offs_n, acc, mask=n_mask) +from freetoken.kernel.triton.bf16_gemv import bf16_gemv def bf16_linear_fp32(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: @@ -46,25 +22,10 @@ def bf16_linear_fp32(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: ``x``: ``[..., K]`` (leading dims collapse to M). ``weight``: ``[N, K]`` bf16. Returns ``[..., N]`` fp32. Bit-exact to ``F.linear(x.float(), weight.float())`` up to fp32 accumulation order; for M>1 it *is* that call (prefill).""" - *lead, K = x.shape - N = weight.shape[0] - M = 1 - for d in lead: - M *= d + M = x.numel() // x.shape[-1] if M != 1: return F.linear(x.float(), weight.float()) - x1 = x.reshape(K).contiguous() - out = torch.empty(N, dtype=torch.float32, device=x.device) - # BN=2 -> N/2 CTAs spreads the (tiny) GEMV across SMs (occupancy-bound, not BW); - # BLOCK_K covering all of K avoids a K-loop. Tuned on H100 (~2 TB/s at N=1024). - BLOCK_N = 2 - BLOCK_K = min(triton.next_power_of_2(K), 4096) - _bf16_gemv_fp32_kernel[(triton.cdiv(N, BLOCK_N),)]( - x1, weight, out, N, K, - weight.stride(0), weight.stride(1), - BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=4, - ) - return out.reshape(*lead, N) + return bf16_gemv(x, weight, torch.float32) __all__ = ["bf16_linear_fp32"] diff --git a/python/freetoken/layers/quantization/linear/unquantized.py b/python/freetoken/layers/quantization/linear/unquantized.py index 6cd8de8a36..3522f465bd 100644 --- a/python/freetoken/layers/quantization/linear/unquantized.py +++ b/python/freetoken/layers/quantization/linear/unquantized.py @@ -1,4 +1,4 @@ -"""bf16 Linear: one kernel (torch), no scheme.""" +"""bf16 Linear: a single bf16 row on CUDA through the Triton GEMV, every other input through torch; no scheme.""" from __future__ import annotations @@ -24,9 +24,24 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: return F.linear(x, w, b) +class TritonLinearKernel(TorchLinearKernel): + """A single bf16 row, decode at batch 1, as one GEMV; every other input is the torch kernel's.""" + + name = "triton" + + def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: + w = layer.weight + single_row = x.numel() == x.shape[-1] + if single_row and x.is_cuda and layer.bias is None and x.dtype == w.dtype == torch.bfloat16: + from freetoken.kernel.triton.bf16_gemv import bf16_gemv + + return bf16_gemv(x, w, x.dtype) + return super().apply(layer, x) + + @register_method(QuantKind.NONE, LayerKind.LINEAR) class UnquantizedLinearMethod(LinearMethod): - candidates = (TorchLinearKernel,) + candidates = (TritonLinearKernel, TorchLinearKernel) def create_weights(self, layer: Any) -> None: g = self.cfg diff --git a/tests/kernels/test_bf16_gemv.py b/tests/kernels/test_bf16_gemv.py new file mode 100644 index 0000000000..bcefaa2f08 --- /dev/null +++ b/tests/kernels/test_bf16_gemv.py @@ -0,0 +1,121 @@ +"""The single-row bf16 GEMV, and the bf16 linear kernel that sends decode's projections through it. + +A decode step at batch 1 multiplies one activation row by every dense weight, so each projection +is a GEMV bound by how fast the weight streams in. The kernel keeps F.linear's contract -- fp32 +accumulation, bf16 or fp32 out -- and the linear kernel keeps F.linear for every other input. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +if not torch.cuda.is_available(): # pragma: no cover + pytest.skip("CUDA required", allow_module_level=True) + +DEV = "cuda" + +# Flash-Next's GDN in_proj, attention qkv, hyper-connection down and up projections; the last +# leaves a tail in both N and K. +SHAPES = [(16480, 2560), (13312, 2560), (336, 10240), (10240, 320), (1001, 6112)] + + +def _weight(N: int, K: int, seed: int) -> torch.Tensor: + torch.manual_seed(seed) + return torch.randn(N, K, device=DEV, dtype=torch.bfloat16) * 0.05 + + +@pytest.mark.parametrize("N,K", SHAPES) +@pytest.mark.parametrize("out_dtype,tolerance", [(torch.bfloat16, 1e-2), (torch.float32, 1e-5)]) +def test_the_gemv_matches_an_fp32_reference(N: int, K: int, out_dtype: torch.dtype, tolerance: float): + from freetoken.kernel.triton.bf16_gemv import bf16_gemv + + weight = _weight(N, K, seed=N + K) + x = torch.randn(1, K, device=DEV, dtype=torch.bfloat16) + + y = bf16_gemv(x, weight, out_dtype) + + reference = x.float() @ weight.float().t() + assert y.shape == (1, N) + assert y.dtype == out_dtype + relative = (y.float() - reference).abs().max() / reference.abs().max() + assert relative.item() < tolerance, relative.item() + + +def test_the_gemv_refuses_a_weight_of_another_width(): + from freetoken.kernel.triton.bf16_gemv import bf16_gemv + + x = torch.randn(1, 2560, device=DEV, dtype=torch.bfloat16) + for weight in (_weight(512, 2048, seed=4), _weight(512, 3072, seed=5)): + with pytest.raises(AssertionError, match="K=2560"): + bf16_gemv(x, weight, torch.bfloat16) + + +def test_the_bf16_linear_method_prefers_the_triton_kernel(): + from freetoken.layers.quantization.linear.base import LinearConfig + from freetoken.layers.quantization.linear.unquantized import UnquantizedLinearMethod + from freetoken.layers.quantization.method import select_kernel + + kernel = select_kernel(UnquantizedLinearMethod.candidates, "auto", LinearConfig(2560, 16480)) + + assert kernel.name == "triton" + + +@pytest.fixture +def gemv_calls(monkeypatch) -> list[tuple[int, ...]]: + """Every shape the linear kernel hands the GEMV, with the GEMV still doing the work.""" + import freetoken.kernel.triton.bf16_gemv as module + + calls: list[tuple[int, ...]] = [] + real = module.bf16_gemv + + def recording(x: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + calls.append(tuple(x.shape)) + return real(x, weight, out_dtype) + + monkeypatch.setattr(module, "bf16_gemv", recording) + return calls + + +@pytest.mark.parametrize( + "rows,activation,with_bias", + [(2, torch.bfloat16, False), (64, torch.bfloat16, False), (1, torch.float32, False), (1, torch.bfloat16, True)], +) +def test_every_other_input_keeps_the_torch_path( + gemv_calls: list[tuple[int, ...]], rows: int, activation: torch.dtype, with_bias: bool +): + from freetoken.layers.quantization.linear.unquantized import TorchLinearKernel, TritonLinearKernel + + weight = _weight(512, 2560, seed=2) + bias = torch.randn(512, device=DEV, dtype=torch.bfloat16) if with_bias else None + layer = SimpleNamespace(weight=weight, bias=bias) + x = torch.randn(rows, 2560, device=DEV, dtype=activation) + + y = TritonLinearKernel().apply(layer, x) + + assert gemv_calls == [] + assert torch.equal(y, TorchLinearKernel().apply(layer, x)) + + +def test_a_single_bf16_row_runs_the_gemv_and_replays_in_a_cuda_graph(gemv_calls: list[tuple[int, ...]]): + from freetoken.layers.quantization.linear.unquantized import TritonLinearKernel + + kernel = TritonLinearKernel() + layer = SimpleNamespace(weight=_weight(13312, 2560, seed=3), bias=None) + x = torch.randn(1, 2560, device=DEV, dtype=torch.bfloat16) + warm = torch.cuda.Stream() + warm.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warm): + kernel.apply(layer, x) + torch.cuda.current_stream().wait_stream(warm) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = kernel.apply(layer, x) + + x.copy_(torch.randn(1, 2560, device=DEV, dtype=torch.bfloat16)) + graph.replay() + + assert torch.equal(y, kernel.apply(layer, x)) + assert gemv_calls == [(1, 2560)] * 3