mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-09-10 20:03:42 +08:00
github-actions[bot]
0e51ecbb7e
version: sunnypilot v2026.003.000 (dev)
date: 2026-09-09T15:43:49
master commit: b255314f5a
19 lines
572 B
Python
19 lines
572 B
Python
from tinygrad import Tensor, dtypes, GlobalCounters
|
|
from tinygrad.dtype import to_dtype
|
|
from tinygrad.helpers import getenv, Context
|
|
from test.backend.test_softmax_fusion import single_kernel_softmax
|
|
|
|
Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
|
|
|
|
if __name__ == "__main__":
|
|
# softmax in bert layers
|
|
BS = getenv("BS", 96//6)
|
|
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
|
|
t = Tensor.empty(BS, 16, 512, 512)
|
|
t.softmax(-1, dtype="half").realize()
|
|
|
|
# test single kernel softmax
|
|
GlobalCounters.reset()
|
|
single_kernel_softmax(t, -1, acc_dtype).realize()
|
|
|