diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index 4f80d4bf2d..b717d5c9fb 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -116,7 +116,7 @@ class TransformerBlock: def _feed_forward(self, h: Tensor) -> Tensor: h_norm = self.ffn_norm(h) - gated = self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm) + gated = (self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)).contiguous(arg=(1,)) return h + self.ffn_down(gated) def __call__(self, x: Tensor, start_pos: int|UOp): diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 68507e9e6e..c8a2a123b7 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -123,7 +123,7 @@ class AddBufferContext: def add_store(ctx:AddBufferContext, x:UOp): rngs = x.src[1:] shape = tuple([r.vmax+1 for r in rngs]) - buf = UOp(Ops.DEFINE_GLOBAL if prod(shape) > 2000 else Ops.DEFINE_LOCAL, dtype=x.dtype.ptr(size=prod(shape)), arg=ctx.dg) + buf = UOp(Ops.DEFINE_GLOBAL if prod(shape) > 65536 else Ops.DEFINE_LOCAL, dtype=x.dtype.ptr(size=prod(shape)), arg=ctx.dg) ctx.map[buf] = (buf.op, ctx.dg) ctx.dg += 1 return buf.reshape(shape).index(*rngs).store(x.src[0], *rngs)