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

remove freqs_cis contiguous in llama (#11597)

Browse Source
This commit is contained in:
chenyu
2025-08-09 21:11:12 -04:00
committed by GitHub
parent 7338ffead0
commit 3e64467322
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/models/llama.py
+2 -3
View File
@@ -181,12 +181,11 @@ class Transformer:
def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
_bsz, seqlen = tokens.shape
h = self.tok_embeddings(tokens)
self.freqs_cis = self.freqs_cis.cast(h.dtype).contiguous()
freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]
mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
# TODO: remove .float()?
logits = self.output(self.norm(h)).float()
if math.isnan(temperature): return logits