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:
@@ -277,6 +277,18 @@ def env_selection(family: str) -> Optional[str]:
|
||||
return env_overrides().get(family)
|
||||
|
||||
|
||||
def env_mode(varname: str) -> str:
|
||||
"""Read a family's ``0``/``1``/``auto`` mode variable (default ``auto``).
|
||||
|
||||
Invalid values warn once per distinct value and fall back to ``auto``.
|
||||
"""
|
||||
mode = os.environ.get(varname, "auto").strip().lower()
|
||||
if mode in ("0", "1", "auto"):
|
||||
return mode
|
||||
_warn_once(f"{varname}={mode!r} is invalid; expected 0, 1, or auto; using auto")
|
||||
return "auto"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Resolution:
|
||||
record: ImplRecord
|
||||
@@ -398,3 +410,30 @@ def explain_plan(calls: Mapping[str, Call]) -> str:
|
||||
return "\n".join(
|
||||
explain(family, *args, **kwargs) for family, (args, kwargs) in calls.items()
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Axes",
|
||||
"Call",
|
||||
"ExplicitSelectionError",
|
||||
"ImplRecord",
|
||||
"OpFamily",
|
||||
"Resolution",
|
||||
"Spec",
|
||||
"Axis",
|
||||
"axis",
|
||||
"env_mode",
|
||||
"env_overrides",
|
||||
"env_selection",
|
||||
"explain",
|
||||
"explain_plan",
|
||||
"get_override",
|
||||
"op_backend",
|
||||
"register_env_alias",
|
||||
"register_family",
|
||||
"reset_override",
|
||||
"resolve",
|
||||
"resolve_plan",
|
||||
"set_override",
|
||||
"tensor_axes",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user