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:

ViperEkura

2026-09-03 07:23:13 +08:00

parent 27abb7c5e7

commit 7540acb43e

14 changed files with 225 additions and 536 deletions

									
										astrai/extension/ops/rotary.py
									
		+3
		
												View File
												
				@@ -29,3 +29,6 @@ def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:

				    if not freqs_cis.is_contiguous():

				        freqs_cis = freqs_cis.contiguous()

				    return mod.rotary_emb(x, freqs_cis)

				__all__ = ["rotary_emb"]