IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.

Skip to main content

For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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.

Structs
Functions