0. Why positional encoding (PE)?Standard attention has no sense of order. Give it “I ate food” or “food ate I” —> same output We need to inject position info explicitly Two flavors: Absolute PE

0. What is Omni-modality? Omni- is a Latin prefix meaning “all” or “every” Omnipotent: All-powerful Omniscient: All-knowing OAI released GPT-4o (GPT-4 Omni), shifting from unimodal (text-only) to o

https://www.youtube.com/watch?v=NFKubflDb1A code was written in 2023 (may be slightly outdated, but the core concept still valid) presented by Sharad Vikram (Pallas author) 1. TPU architecture recap

1. jax.jit jax.jit traces Python into computational graph (a jaxpr) → XLA compiles the graph into an optimized HLO program for the target device After compilation, Python is completely out of the loo

1. Longformerhttps://arxiv.org/abs/2004.05150 it was the 1st sliding-window attention (SWA) paper. published in 2020 traditional full attention: compute complexity $O(n^2)$, memory $O(n^2)$, where

https://openxla.org/xla/shapes https://openxla.org/xla/tiled_layout 1. XLA op formatHLO example 12add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)} add(exponential.183, broadcas

https://openxla.org/xla 0. IntroXLA in the whole JAX stack Source: Yi Wang’s linkedin post LLM is basically matmul. XLA (Accelerated Linear Algebra) optimizes linear algebra on multiple decives (TPU

source: https://huggingface.co/blog/Kseniase/kd 1. history Knowledge Distillation (KD): transfer knowledge from teacher model to a smaller student model DeepSeek-R1 proposed effective distillation imp

https://www.youtube.com/watch?v=NQ-0D5Ti2dc&t=9s https://github.com/gpu-mode/lectures/tree/main/lecture_002 from PMPP book 1. Memory allocation nvidia devices come with their own DRAM (device) glo

1. why pallas? JAX works with pure functions (i.e., same inputs will produce the same outputs). JAX arrays are immutable not flexible or efficient for kernel implementation GEMM steps input matrix →