Intro to Pallas blog post #3
Open
+212
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sharadmv
reviewed
Nov 19, 2025
sharadmv
reviewed
Nov 19, 2025
sharadmv
reviewed
Nov 19, 2025
MichalPitr
reviewed
Nov 19, 2025
jakevdp
reviewed
Nov 19, 2025
justinjfu
force-pushed
the
intro_pallas
branch
from
5061c6d to
f9be0d0
Compare
justinjfu
force-pushed
the
intro_pallas
branch
from
f9be0d0 to
d2dcdfc
Compare
superbobry
approved these changes
Nov 21, 2025
|
|
||
| However, the stated FLOP/s does not tell the whole story as not all operations can be performed at the same speed. The vast majority (often 90%+) of the peak FLOP/s on modern accelerators is restricted to matrix multiplications performed by the TensorCore (GPUs) or MXU (TPUs). Element-wise operations such as point-wise addition or multiplication are performed at a much slower rate by the vector ALU on the chip. Scalar FLOP/s are also much lower, especially compared to CPUs. This means that if your workflow is not dominated by matrix multiplications, you should never expect to utilize more than 10-20% of the FLOP/s of your accelerator. | ||
|
|
||
|  |
There was a problem hiding this comment.
Is there a diagram for v5e? or can we say that the high level design for v5e is similar? We are talking about v5e in this section and then show a diagram for a different generation.
| Modern accelerators (such as TPUs and GPUs) typically contain multiple levels of memory just like CPUs, each trading off speed vs. capacity. | ||
|
|
||
| - **Registers** live closest to the processor and can usually be accessed within a single processor cycle. Each core of an accelerator typically has a very limited number of registers (typically in the 10-100 range, depending on the type of register). | ||
| - **Caches** (SRAM) typically live on same die as the processor, and can typically be accessed with 1-10x the latency of accessing a register. The amount of SRAM usually ranges in the 10s-100s MB range. |
There was a problem hiding this comment.
I wonder if it is accurate to call SRAM "caches"? Is this how it's used on TPU?
| 2. Here, we need to explicitly zero-out `out_vmem_ref` on the first accumulation step | ||
| of each output block as Pallas does not guarantee any initialization for buffers. | ||
|
|
||
| We construct a Pallas kernel using the `pl.kernel` entry point. Passing in `mesh=pltpu.create_tensorcore_mesh("x")` informs JAX that we wish to use the TPU tensorcore backend (as opposed to GPU or TPU sparsecore). On TPUs with megacore enabled, the axis name we passed in can be queried to obtain the core id via `lax.axis_index("x")`. |
There was a problem hiding this comment.
Nit: should we capitalize first letters in tensorcore and sparsecore?
| 2. Here, we need to explicitly zero-out `out_vmem_ref` on the first accumulation step | ||
| of each output block as Pallas does not guarantee any initialization for buffers. | ||
|
|
||
| We construct a Pallas kernel using the `pl.kernel` entry point. Passing in `mesh=pltpu.create_tensorcore_mesh("x")` informs JAX that we wish to use the TPU tensorcore backend (as opposed to GPU or TPU sparsecore). On TPUs with megacore enabled, the axis name we passed in can be queried to obtain the core id via `lax.axis_index("x")`. |
There was a problem hiding this comment.
Nit: using something like "core" instead of "x" might be slightly more clear. "x" also suggests that the axis might be use for partitioning the x input, which isn't the case in the kernel.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.