nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

fix mxfp4 mem estimate (#17515)

Browse Source 
* add mem estimates

* rename

* move
This commit is contained in:
qazal
2026-08-13 14:48:15 +09:00
committed by GitHub
parent e103fb2a10
commit 39d144546e
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/cdna_asm_gemm.py
+2 -1
View File
@@ -122,7 +122,8 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
insts = build_kernel(M, N, K, tile_m, tile_n)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))