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:
@@ -83,7 +83,10 @@ def test_mode_one_forces_supported_shape(monkeypatch):
|
||||
|
||||
|
||||
@skip_no_swiglu
|
||||
def test_auto_falls_back_until_shape_is_qualified(monkeypatch):
|
||||
def test_auto_uses_unfused_chain_until_shape_is_qualified(monkeypatch):
|
||||
# The fusion table is empty, so auto keeps the unfused linear-backend
|
||||
# chain. The linear backend may still dispatch its own GEMV for M=4,
|
||||
# hence the relaxed tolerance versus the pure-torch reference.
|
||||
monkeypatch.setenv("ASTRAI_SWIGLU", "auto")
|
||||
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
@@ -91,4 +94,4 @@ def test_auto_falls_back_until_shape_is_qualified(monkeypatch):
|
||||
with torch.no_grad():
|
||||
actual = swiglu(x, up_weight, gate_weight)
|
||||
expected = reference_swiglu(x, up_weight, gate_weight)
|
||||
torch.testing.assert_close(actual, expected)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.1)
|
||||
|
||||
Reference in New Issue
Block a user