For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo module
mla_prefill_per_token_scale
Per-token-scale MLA prefill kernel.
Thin dispatch wrapper that creates scale TMA tiles and calls the generic MLA kernel's per-token-scale variant (mla_prefill_kernel_per_token_scale).
Structs
-
MLASmemStorage: Shared memory storage layout for the per-token-scale MLA prefill kernel.
Functions
-
mla_sm100_prefill_per_token_scale: Host-side entry point for the SM100 MLA prefill kernel with per-token scaling. -
q_scale_tma: Creates a 2-D TMA tile descriptor for the per-token Q scale tensor.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!