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:
2026-09-03 07:23:13 +08:00
parent 27abb7c5e7
commit 7540acb43e
14 changed files with 225 additions and 536 deletions
+39
View File
@@ -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",
]