beam: add dev_timeout for am (#15063)

* beam: add dev_timeout for am

* all covered

* fk

* x

* fuzz

* reset

* f
This commit is contained in:
nimlgen
2026-03-01 16:57:29 +03:00
committed by GitHub
parent 8cb4368967
commit ccbbca05ef
18 changed files with 79 additions and 41 deletions
+1 -1
View File
@@ -41,7 +41,7 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_
class PythonProgram:
def __init__(self, name:str, lib:bytes, **kwargs):
self.uops: list[tuple[Ops, DType, list[int], Any]] = pickle.loads(lib)
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
st = time.perf_counter()
warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
warp_size = len(warp)