Revert "Add .warmup() for triton.jit()" (#682)
Reverts openai/triton#671 It seems like for some reason this caused out-of-memory errors on some of our internal workloads. I'm reverting this so that HEAD can be used in production at OpenAI, and I will work on digging into this issue asynchronously.
This commit is contained in:
@@ -150,25 +150,3 @@ def test_constexpr_not_callable() -> None:
|
||||
except BaseException:
|
||||
error = True
|
||||
assert error is True
|
||||
|
||||
|
||||
def test_jit_warmup_cache() -> None:
|
||||
@triton.jit
|
||||
def kernel_add(a, b, o, N: tl.constexpr):
|
||||
idx = tl.arange(0, N)
|
||||
tl.store(o + idx,
|
||||
tl.load(a + idx) + tl.load(b + idx))
|
||||
|
||||
args = [
|
||||
torch.randn(32, dtype=torch.float32, device="cuda"),
|
||||
torch.randn(32, dtype=torch.float32, device="cuda"),
|
||||
torch.randn(32, dtype=torch.float32, device="cuda"),
|
||||
32,
|
||||
]
|
||||
assert len(kernel_add.cache) == 0
|
||||
kernel_add[(1,)].warmup(torch.float32, torch.float32, torch.float32, 32)
|
||||
assert len(kernel_add.cache) == 1
|
||||
kernel_add[(1,)].warmup(*args)
|
||||
assert len(kernel_add.cache) == 1
|
||||
kernel_add[(1,)](*args)
|
||||
assert len(kernel_add.cache) == 1
|
||||
|
Reference in New Issue
Block a user