perf: dispatch linear gemv by decode batch size and unify extension style
- 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].
This commit is contained in:
@@ -9,6 +9,8 @@ Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
|
||||
freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
@@ -41,7 +43,7 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
||||
return x_out.to(dtype)
|
||||
|
||||
|
||||
def _rotary_records() -> list:
|
||||
def _rotary_records() -> List[ImplRecord]:
|
||||
return [
|
||||
ImplRecord(
|
||||
family="rotary",
|
||||
@@ -61,12 +63,15 @@ def _rotary_records() -> list:
|
||||
]
|
||||
|
||||
|
||||
register_family(
|
||||
"rotary",
|
||||
lambda x, freqs_cis: tensor_axes(x),
|
||||
_rotary_records,
|
||||
lambda: _rotary_records()[-1],
|
||||
)
|
||||
def _axes(x: Tensor, freqs_cis: Tensor) -> Dict[str, Any]:
|
||||
return tensor_axes(x)
|
||||
|
||||
|
||||
def _fallback_record() -> ImplRecord:
|
||||
return _rotary_records()[-1]
|
||||
|
||||
|
||||
register_family("rotary", _axes, _rotary_records, _fallback_record)
|
||||
|
||||
|
||||
def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
||||
@@ -80,3 +85,6 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
||||
[batch, seq_len, n_heads, head_dim] (bf16)
|
||||
"""
|
||||
return resolve("rotary", x, freqs_cis).record.obj(x, freqs_cis)
|
||||
|
||||
|
||||
__all__ = ["apply_rotary_emb"]
|
||||
|
||||
Reference in New Issue
Block a user