text
Triton nanoGPT now has a custom cross entropy loss kernel 🚀 
Triton nanoGPT now has a custom cross entropy loss kernel 🚀 
new_line
text
Next: matmul, gradually overthrowing all major PyTorch ops:) 
Next: matmul, gradually overthrowing all major PyTorch ops:) 
new_line
new_line
text
Simplified pseudo for parallel cross-entropy loss compute:
Simplified pseudo for parallel cross-entropy loss compute:
new_line
text
- init program: get pid, compute offsets, load targets.
- init program: get pid, compute offsets, load targets.
new_line
text
- init row_max and row_sum.
- init row_max and row_sum.
new_line
text
- for-loop1 (find max logits): update row_max with max logits.
- for-loop1 (find max logits): update row_max with max logits.
new_line
text
- for-loop2 (compute softmax and loss): compute row_sum, update loss.
- for-loop2 (compute softmax and loss): compute row_sum, update loss.
new_line
text
- add log(row_sum) and store loss. 
- add log(row_sum) and store loss. 
new_line
new_line
text
Code: 
Code: 
link
https://github.com/Jaykef/ai-algorithms/blob/main/triton_nanoGPT.ipynb
https://github.com/Jaykef/ai-algorithms/blob/main/triton_nanoGPT.ipynb
Triton nanoGPT now has a custom cross entropy loss kernel 🚀 
Next: matmul, gradually overthrowing all major PyTorch ops:) 

Simplified pseudo for parallel cross-entropy loss compute:
- init program: get pid, compute offsets, load targets.
- init row_max and row_sum.
- for-loop1 (find max logits): update row_max with max logits.
- for-loop2 (compute softmax and loss): compute row_sum, update loss.
- add log(row_sum) and store loss. 

Code: https://github.com/Jaykef/ai-algorithms/blob/main/triton_nanoGPT.ipynb
🔥
Sri-Vigneshwar-DJ John6666 sugatoray Joseph717171
👀
John6666 Joseph717171