perf(linear): run a single-row bf16 projection as one GEMV

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"]

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"]

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

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