perf: accelerate decode linear with bf16 gemv
- add decode-shape benchmark harness - add bf16 GEMV CUDA primitive with head-dim generic kernel - dispatch decode-time linear layers to gemv for M=1 - extend gemv coverage to small decode batches
This commit is contained in:
0z5a
@@ -0,0 +1,155 @@
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import bf16_gemv, is_available
|
||||
|
||||
GEMV_AVAILABLE = (
|
||||
torch.cuda.is_available()
|
||||
and is_available("bf16_gemv")
|
||||
and torch.cuda.get_device_capability() >= (8, 0)
|
||||
)
|
||||
skip_no_gemv = pytest.mark.skipif(
|
||||
not GEMV_AVAILABLE,
|
||||
reason="BF16 GEMV requires a built kernel and compute capability 8.0+",
|
||||
)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
@pytest.mark.parametrize(
|
||||
"n,k",
|
||||
[(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)],
|
||||
)
|
||||
def test_bf16_gemv_matches_linear_shape_families(n, k):
|
||||
torch.manual_seed(17)
|
||||
x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
|
||||
actual = bf16_gemv(x, weight)
|
||||
expected = F.linear(x, weight)
|
||||
assert actual.shape == (n,)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
@pytest.mark.parametrize("m", [2, 4, 8])
|
||||
@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)])
|
||||
def test_bf16_gemv_matches_small_decode_batches(m, n, k):
|
||||
torch.manual_seed(19 + m)
|
||||
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
|
||||
actual = bf16_gemv(x, weight)
|
||||
expected = F.linear(x, weight)
|
||||
assert actual.shape == (m, n)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_bf16_gemv_preserves_singleton_batch_and_fuses_bias():
|
||||
torch.manual_seed(23)
|
||||
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
bias = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
|
||||
actual = bf16_gemv(x, weight, bias)
|
||||
expected = F.linear(x, weight, bias)
|
||||
assert actual.shape == (1, 1536)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_bf16_gemv_small_batch_fuses_bias():
|
||||
torch.manual_seed(25)
|
||||
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
|
||||
actual = bf16_gemv(x, weight, bias)
|
||||
expected = F.linear(x, weight, bias)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_bf16_gemv_uses_current_stream():
|
||||
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
stream = torch.cuda.Stream()
|
||||
with torch.cuda.stream(stream):
|
||||
actual = bf16_gemv(x, weight)
|
||||
expected = F.linear(x, weight)
|
||||
stream.synchronize()
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_bf16_gemv_cuda_graph_replay():
|
||||
torch.manual_seed(29)
|
||||
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
for _ in range(3):
|
||||
bf16_gemv(x, weight)
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
actual = bf16_gemv(x, weight)
|
||||
|
||||
x.copy_(torch.randn_like(x))
|
||||
graph.replay()
|
||||
expected = F.linear(x, weight)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_bf16_gemv_small_batch_cuda_graph_replay():
|
||||
torch.manual_seed(31)
|
||||
x = torch.randn(8, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
for _ in range(3):
|
||||
bf16_gemv(x, weight)
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
actual = bf16_gemv(x, weight)
|
||||
|
||||
x.copy_(torch.randn_like(x))
|
||||
graph.replay()
|
||||
expected = F.linear(x, weight)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
@pytest.mark.parametrize(
|
||||
"make_args,error",
|
||||
[
|
||||
(
|
||||
lambda: (
|
||||
torch.randn(3, 16, device="cuda", dtype=torch.bfloat16),
|
||||
torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
|
||||
),
|
||||
"M must",
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
torch.randn(15, device="cuda", dtype=torch.bfloat16),
|
||||
torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
|
||||
),
|
||||
"even",
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
torch.randn(16, device="cuda", dtype=torch.float16),
|
||||
torch.randn(8, 16, device="cuda", dtype=torch.float16),
|
||||
),
|
||||
"bf16",
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
torch.randn(
|
||||
16, device="cuda", dtype=torch.bfloat16, requires_grad=True
|
||||
),
|
||||
torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
|
||||
),
|
||||
"autograd",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_bf16_gemv_rejects_unsupported_inputs(make_args, error):
|
||||
with pytest.raises(RuntimeError, match=error):
|
||||
bf16_gemv(*make_args())
|
||||
@@ -0,0 +1,185 @@
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import explain, is_available, linear, op_backend
|
||||
from astrai.extension.backend import linear as public_linear
|
||||
from astrai.model.components.linear import Linear
|
||||
|
||||
GEMV_AVAILABLE = (
|
||||
torch.cuda.is_available()
|
||||
and is_available("bf16_gemv")
|
||||
and torch.cuda.get_device_capability() >= (8, 0)
|
||||
)
|
||||
skip_no_gemv = pytest.mark.skipif(
|
||||
not GEMV_AVAILABLE,
|
||||
reason="BF16 GEMV requires a built kernel and compute capability 8.0+",
|
||||
)
|
||||
|
||||
|
||||
def test_linear_backend_is_public():
|
||||
assert linear is public_linear
|
||||
|
||||
|
||||
def test_model_linear_routes_through_backend(monkeypatch):
|
||||
sentinel = torch.randn(2, 4)
|
||||
|
||||
def fake_linear(x, weight, bias):
|
||||
assert x.shape == (2, 3)
|
||||
assert weight.shape == (4, 3)
|
||||
assert bias is None
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr("astrai.model.components.linear.linear", fake_linear)
|
||||
layer = Linear(3, 4)
|
||||
assert layer(torch.randn(2, 3)) is sentinel
|
||||
|
||||
|
||||
def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "1")
|
||||
x = torch.randn(2, 8, requires_grad=True)
|
||||
weight = torch.randn(4, 8, requires_grad=True)
|
||||
actual = linear(x, weight)
|
||||
expected = F.linear(x, weight)
|
||||
torch.testing.assert_close(actual, expected)
|
||||
actual.sum().backward()
|
||||
assert x.grad is not None
|
||||
assert weight.grad is not None
|
||||
|
||||
|
||||
def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "invalid-test-mode")
|
||||
with caplog.at_level(logging.WARNING):
|
||||
trace = explain("linear", torch.randn(1, 8), torch.randn(4, 8))
|
||||
assert "using auto" in caplog.text
|
||||
assert "=> torch" in trace
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_mode_zero_disables_gemv(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "0")
|
||||
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
assert "=> torch" in explain("linear", x, weight)
|
||||
torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_mode_one_forces_capable_unmeasured_shape(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "1")
|
||||
x = torch.randn(1, 64, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
assert "=> gemv" in explain("linear", x, weight)
|
||||
torch.testing.assert_close(
|
||||
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
|
||||
)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
@pytest.mark.parametrize("m", [2, 4, 8])
|
||||
def test_mode_one_dispatches_supported_small_batches(monkeypatch, m):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "1")
|
||||
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
assert "=> gemv" in explain("linear", x, weight)
|
||||
torch.testing.assert_close(
|
||||
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
|
||||
)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "auto")
|
||||
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
winning = torch.randn(
|
||||
1536,
|
||||
1536,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
requires_grad=True,
|
||||
)
|
||||
with torch.no_grad():
|
||||
assert "=> torch" in explain("linear", x, winning)
|
||||
torch.testing.assert_close(linear(x, winning), F.linear(x, winning))
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "auto")
|
||||
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
trace = explain("linear", x, weight)
|
||||
if torch.cuda.get_device_capability() == (8, 9):
|
||||
assert "=> auto_gemv" in trace
|
||||
else:
|
||||
assert "=> torch" in trace
|
||||
torch.testing.assert_close(
|
||||
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.5
|
||||
)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
@pytest.mark.parametrize(
|
||||
"m,n,k",
|
||||
[
|
||||
(2, 6912, 1536), # up/gate loses at every measured M
|
||||
(2, 1536, 6912), # long-K accumulation changed checkpoint greedy output
|
||||
(4, 100000, 1536), # LM head misses the 5% M=4 gate
|
||||
(4, 1536, 6912), # long-K accumulation changed checkpoint greedy output
|
||||
(8, 256, 1536), # remaining M=8 winners miss the 3% end-to-end gate
|
||||
],
|
||||
)
|
||||
def test_auto_rejects_measured_small_batch_losers(monkeypatch, m, n, k):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "auto")
|
||||
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
assert "=> torch" in explain("linear", x, weight)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_grad_enabled_and_unsupported_multirow_always_fall_back(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "1")
|
||||
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(
|
||||
256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
|
||||
)
|
||||
assert "=> torch" in explain("linear", x, weight)
|
||||
with torch.no_grad():
|
||||
multirow = x.expand(3, -1).contiguous()
|
||||
assert "=> torch" in explain("linear", multirow, weight)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_explicit_gemv_selection_respects_capability(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "0")
|
||||
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad(), op_backend(linear="gemv"):
|
||||
assert "=> gemv" in explain("linear", x, weight)
|
||||
torch.testing.assert_close(
|
||||
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
|
||||
)
|
||||
|
||||
|
||||
@skip_no_gemv
|
||||
def test_dispatched_linear_cuda_graph_replay(monkeypatch):
|
||||
monkeypatch.setenv("ASTRAI_GEMV", "1")
|
||||
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
for _ in range(3):
|
||||
linear(x, weight)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
actual = linear(x, weight)
|
||||
x.copy_(torch.randn_like(x))
|
||||
graph.replay()
|
||||
expected = F.linear(x, weight)
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
|
||||
Reference in New Issue
Block a user