ViperEkura
7540acb43e
- replace the per-shape auto tables in the linear backend with an M-banded rule (M in [2,4] on compute capability 8.0+) that measured at the HBM bandwidth floor across every family, and fold the capability check into the capable guard - drop the unreachable swiglu auto shape-table machinery so both backends share one env-mode ladder via the new dispatch.env_mode helper - add __all__ across extension modules, name the rotary registration records, and unify typing to the typing-module style - rewrite test_linear_dispatch.py around behavioral routing assertions and document the M-banded policy in the developer docs - Benchmark: L20 SM89, Python dispatch overhead 2.9us to 1.5us, auto now covers every projection shape at M in [2,4].
27 lines
791 B
Python
27 lines
791 B
Python
"""Stateless wrapper for the directly callable BF16 GEMV primitive."""
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from astrai.extension.loader import get_module
|
|
|
|
|
|
def bf16_gemv(
|
|
x: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
"""Compute ``F.linear(x, weight, bias)`` for up to eight BF16 rows.
|
|
|
|
``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 8]``,
|
|
and ``weight`` must be a contiguous row-major ``[N, K]`` tensor. The CUDA
|
|
kernel reuses each weight row across M, accumulates in FP32, and returns
|
|
BF16. This primitive is inference-only and intentionally performs no
|
|
fallback or model-level dispatch.
|
|
"""
|
|
return get_module("bf16_gemv").bf16_gemv(x, weight, bias)
|
|
|
|
|
|
__all__ = ["bf16_gemv"]
|