tform ffn

This commit is contained in:
2025-08-08 18:29:48 -07:00
parent ee06481036
commit e268eb2d5c
2 changed files with 2 additions and 2 deletions
+1 -1
View File
@@ -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):
+1 -1
View File
@@ -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)