diff --git a/README.md b/README.md index aaedf18e..70225611 100644 --- a/README.md +++ b/README.md @@ -23,8 +23,9 @@ vehicle is **TIRx** (Tensor IR next), a Python DSL for writing GPU kernels at th layout model (`TileLayout`, named axes, swizzle). - **Part III — GEMM: Tiled to SOTA.** A tiled GEMM built up through TMA pipelining, persistent scheduling, warp specialization, and 2-CTA clusters. -- **Part IV — Flash Attention 4.** A complete attention kernel built from the Part III techniques: - two MMAs with softmax between them, online-softmax rescaling, causal masking, and GQA. +- **Part IV — Attention Kernels.** Flash Attention 4 and FlashMLA connect the Part III techniques + into complete attention pipelines, including online softmax, latent KV representations, sparse + gathers, warp specialization, and cooperative two-CTA execution. - **Appendices.** TIRx language reference, [reproducible GPU benchmarking and profiling](appendix/benchmarking_gpu_kernels.md), compiler internals, and asynchronous-kernel debugging. diff --git a/chapter_flashmla/index.md b/chapter_flashmla/index.md new file mode 100644 index 00000000..a8190db5 --- /dev/null +++ b/chapter_flashmla/index.md @@ -0,0 +1,1366 @@ +(chap_flashmla)= +# FlashMLA + +:::{admonition} Overview +:class: overview + +- Start from an ordinary MHA KV cache and derive why MLA stores one shared + compressed state per token, including where the head-specific K/V + transformations move. +- Define FlashMLA sparse attention: an external indexer selects KV rows, then + the kernel performs QK, softmax, and PV over those rows. An executable + reference fixes its numerical behavior and edge cases. +- Use a concrete implementation to understand how sparse attention maps to + Blackwell, then compile and verify it on B200. +::: + +A language model normally predicts one new token at a time. At every step, +attention compares the query at the current position with the keys and values +of all earlier tokens, then uses the attention weights to gather information +from that history. Without saving the historical K/V states, the model would +recompute K/V for the entire prefix at every step. A **KV cache** avoids that +repeated work by storing those states once and letting later attention steps +read them directly. Generative models normally use causal attention, so the K/V +at position $s$ depend only on tokens up to position $s$ and remain unchanged +when later tokens are appended. Prefill computes the prompt tokens in parallel +and fills the cache; decode then reads that cache step by step and appends the +new token's K/V. + +In ordinary multi-head attention (MHA), every attention head has its own K and +V, so the KV cache must store every head's K/V for every processed token. As the +context grows, the cache consumes more GPU memory, and attention must read an +ever-longer K/V history for every newly generated token. It can therefore +become both a capacity and a memory-bandwidth bottleneck. **Multi-head Latent +Attention (MLA)** reduces these costs by changing the cache representation: it +compresses the K/V-related information for each token into one state shared by +all heads. During attention, every head still applies its own transformations +and therefore retains distinct attention behavior. + +FlashMLA is DeepSeek's library of optimized GPU kernels for MLA. This chapter +studies one sparse-attention forward kernel for Blackwell that is used during +prefill. While the prompt tokens are processed in parallel, an external indexer +first selects the history positions that each query token should attend to. The +kernel loads the corresponding KV rows, then performs QK, softmax, and PV. Compared with +dense attention, it reduces only the history positions that participate in the +calculation; the main attention chain remains unchanged. + +In the kernel studied here, each query token corresponds to 128 query heads, +while every selected history token has one KV-cache row shared by those heads. +But if 128 heads read the same row, how can they still produce different +results? The cached compressed state is shared, but every head retains its own +transformations on the query and output paths. To see how those transformations +move, start with the KV cache of ordinary MHA. + +## From the ordinary MHA KV cache to MLA + +### Ordinary MHA KV-cache cost + +Ordinary MHA has $n_h$ heads of width $d_h$, each with its own query, key, and +value projections. Fix one head and write its projection matrices as $W^Q$, +$W^K$, and $W^V$. Within one layer, write the current token vector entering the +attention projections as $h_t$ and the input vector for the earlier token at +position $s$ as $h_s$. Then + +$$ +q=W^Q h_t,\qquad +k_s=W^K h_s,\qquad +v_s=W^V h_s. +$$ + +This head computes + +$$ +p_s=\operatorname{softmax}_s +\left(\frac{q^{\mathsf T}k_s}{\sqrt{d_h}}\right), +\qquad +o=\sum_s p_s v_s. +$$ + +During autoregressive generation, the current query can be discarded after use, +but later queries repeatedly read the historical $k_s$ and $v_s$. +Each layer therefore caches $2n_hd_h$ elements per token: one K and one V for +every head. + +The total grows linearly with batch size, context length, and layer count. For +example, 32 layers, context length 4096, 32 heads of width 128, batch size 1, +and BF16 require about 2 GB of KV cache. Generating every new token +also rereads this growing history, so KV cache creates both memory-capacity and +bandwidth pressure. + +One direct way to reduce this cost is to share more state across query heads. +Multi-query attention (MQA) shares one K/V pair across all query heads, while +grouped-query attention (GQA) shares within groups. MLA uses a different form of +sharing: every head retains its own projection, but the cache stores only one +shared low-dimensional state. + +### From a shared state to a low-dimensional latent cache + +Start with the non-positional content path and write the current query content +as $q^C$. Ordinary MHA first forms $k_s^C=W^Kh_s$, then computes the content +score in QK. Associativity lets us rewrite that step as + +$$ +\mathrm{QK}_s +=(q^{C})^{\mathsf T}k_s^C +=(q^{C})^{\mathsf T}W^Kh_s +=\left((W^K)^{\mathsf T}q^{C}\right)^{\mathsf T}h_s. +$$ + +PV likewise forms $v_s=W^Vh_s$ before taking the attention-weighted sum with +weights $p_s$. Linearity gives + +$$ +\mathrm{PV} +=\sum_s p_s v_s +=\sum_s p_s W^Vh_s +=W^V\left(\sum_s p_s h_s\right). +$$ + +The first identity applies $(W^K)^{\mathsf T}$ to the current query once and +then dots the result with every $h_s$. The second first takes the weighted sum +of the $h_s$ vectors and only then applies $W^V$. The numerical results stay the +same, while the cache only needs to retain one $h_s$. Think of this as an +uncompressed latent-cache thought experiment. + +$h_s$ is still $d_{model}$ coordinates wide. Caching it shares state across +heads, but QK scoring and PV aggregation still operate in that wide space. +Actual MLA uses two steps. + +First, every head shares a matrix $D$ that compresses $h_s$ into a +$d_c$-dimensional vector $c_s$. For the content path, the KV cache stores only +this $c_s$: + +$$ +c_s=Dh_s,\qquad c_s\in\mathbb{R}^{d_c}. +$$ + +Second, each head has its own up-projection matrices. Continue to fix one head +and write its matrix blocks as $U_K$ and $U_V$. If K/V are expanded explicitly, +this head's content key and value are + +$$ +k_s^C=U_Kc_s,\qquad +v_s=U_Vc_s. +$$ + +These equations define the K/V that this head must use. Weight absorption later +preserves the same result without materializing these intermediate vectors. +$D$, $U_K$, and $U_V$ are learned together during model training. Because $d_c$ +is much smaller than $d_{model}$, the content state cached for each token is +correspondingly smaller. + +### A separate Q/K feature block for RoPE + +For a content score without positional information, the current query can be +multiplied by $U_K^{\mathsf T}$ once and the result reused for every historical +position. After adding RoPE, we first need to check whether that reuse remains +possible. + +Let $q_t^C$ be the current query's content feature, and let $R_t$ and $R_s$ be +the RoPE rotations at query position $t$ and historical position $s$. Applying +RoPE directly to the content query and content key would give + +$$ +\left(R_tq_t^{C}\right)^{\mathsf T} +\left(R_sU_Kc_s\right) += +\left(U_K^{\mathsf T}R_s^{\mathsf T}R_tq_t^{C}\right)^{\mathsf T}c_s. +$$ + +The query position $t$ is fixed, but the key position $s$ changes as the query +is compared with different historical tokens. The vector +$U_K^{\mathsf T}R_s^{\mathsf T}R_tq_t^C$ therefore also changes with $s$ and +cannot remain a single vector determined only by the current query. Once RoPE +is applied directly to content Q/K, the query-side result that was meant to be +reused now changes with each historical position. + +MLA therefore produces a separate, narrower block of Q/K features and applies +RoPE only to that block. Let $q_t^R$ and $k_s^R$ denote the results after the +RoPE rotations at positions $t$ and $s$. The superscript $R$ means that these +features have passed through RoPE; they are still learned projections of token +representations rather than pure position vectors. $q_t^R$ belongs to the +current head, while $k_s^R$ is shared by all heads. + +In the expanded-K/V form, the complete query and key are +$[q_t^C;q_t^R]$ and $[U_Kc_s;k_s^R]$, where the semicolon denotes +concatenation along the feature dimension. Their complete QK dot product is + +$$ +\mathrm{score}_s +=[q_t^C;q_t^R]^{\mathsf T}[U_Kc_s;k_s^R] +=(q_t^C)^{\mathsf T}U_Kc_s + +(q_t^R)^{\mathsf T}k_s^R. +$$ + +The sum combines two scalar contributions from different coordinate blocks of +the same QK dot product; it does not add the two vectors. The first term +measures content similarity. The second matches RoPE-transformed Q/K features, +adding relative-position information to the score. + +### Weight absorption: avoiding per-head K/V expansion with associativity + +Evaluating these formulas directly would multiply every cached $c_s$ by $U_K$ +and $U_V$, explicitly producing this head's historical content K/V. Weight +absorption is the use of associativity to change the parenthesization and avoid +those intermediate tensors. + +First consider the content term in the preceding score. The key side can be +rewritten as + +$$ +(q_t^C)^{\mathsf T}(U_Kc_s) +=(U_K^{\mathsf T}q_t^C)^{\mathsf T}c_s. +$$ + +The left side forms the content key $U_Kc_s$ at every historical position. The +right side computes $U_K^{\mathsf T}q_t^C$ once for the current query and dots +the result with every cached $c_s$. Replacing the first term of the preceding +score with the right side leaves the RoPE term and the complete score unchanged. +Using the same scale, mask, and softmax therefore produces the same attention +weights $p_s$, without materializing $U_Kc_s$ as a per-head content key. + +The value side permits the same change of evaluation order. Let $p_s$ be this +head's attention weight for historical position $s$. Then + +$$ +o +=\sum_s p_s v_s +=\sum_s p_s U_Vc_s +=U_V\left(\sum_s p_s c_s\right). +$$ + +The original order first forms $v_s=U_Vc_s$ for every historical position and +then takes their weighted sum. After regrouping, attention first computes +$\sum_s p_s c_s$ in the shared latent space and applies $U_V$ only once to the +result. Multi-head attention finally uses the output projection $W_O$ to +combine the heads; because $U_V$ is also linear, it can be composed in advance +with this head's matrix block in $W_O$. + +Both the key-side and value-side changes only regroup the multiplications; the +model parameters and numerical result remain unchanged. Together these two +regroupings are called weight absorption. The same MLA layer therefore has two +equivalent evaluation modes: explicitly producing each head's K/V is called +MHA mode, while evaluating attention directly on the shared latent state with +the regrouping above is called MQA mode. + +```{figure} ../img/flashmla_mla_modes.png +:width: 100% +:alt: MHA and MQA execution modes of MLA and key-side and value-side weight absorption + +*MHA mode explicitly expands each head's K/V. MQA mode computes +$U_K^{\mathsf T}q_t^C$ before QK and applies $U_V$ after the latent weighted +sum; both modes use the same RoPE term and produce the same result.* +``` + +The two modes produce the same result but incur different costs. MHA mode must +first produce each head's K/V, but QK/PV use a narrower feature width; the +expansion can be amortized when many queries reuse it. MQA mode removes those +intermediate K/V tensors, but QK and PV work directly in the wider latent +space. The right choice depends on the usage stage, sparsity, tensor shape, +data movement, and hardware schedule. + +The sparse-prefill kernel studied in this chapter uses MQA mode. + +### MLA cache contents and footprint + +In MQA mode, each historical token caches only the shared $c_s$ and the shared +RoPE-transformed key feature $k_s^R$, rather than expanded per-head content K/V. +Every head still has its own query, $U_K$, and $U_V$, so heads can produce +different attention weights and outputs. The following figure compares the +ordinary MHA cache with the shared MLA cache: + +```{figure} ../img/flashmla_cache_story.png +:width: 100% +:alt: Ordinary MHA caches separate key and value data for every head; MLA stores one shared compressed state and keeps head-specific work around attention + +*Ordinary MHA stores a separate key/value slice per head. MLA stores one shared +compressed content state and one shared RoPE-transformed key feature per token; +head-specific query and output transformations happen before and after +attention.* +``` + +We can now compare what each mechanism actually caches. Here $n_{kv}$ is the +number of GQA KV heads, and $d_h^R$ is the width of $k_s^R$. “Ordinary MQA” in +the table names an attention architecture rather than MLA's MQA execution mode. +The counts are scalar elements stored per token per layer, excluding dtype, +alignment, and other overhead: + +| Mechanism | Cached elements | Cached state | +| --- | ---: | --- | +| MHA | $2n_hd_h$ | Complete K and V for every head | +| GQA | $2n_{kv}d_h$ | Complete K and V for every KV head | +| Ordinary MQA | $2d_h$ | One complete K/V pair shared by all query heads | +| MLA | $d_c+d_h^R$ | One $c_s$ and $k_s^R$ shared by all heads | + +MLA does not cache $2d_c$ elements: the same $c_s$ contains the information +needed to form both content K and V and is stored only once. + +This chapter uses $d_c=512$ and $d_h^R=64$, so every $[c_s;k_s^R]$ contains +$512+64=576$ scalar elements. The later kernel parameter `d_qk=576` denotes +the width of this cached row. + +### Numerical verification of the two execution modes + +The CPU program constructs both executions. It includes a shared RoPE score +term, expands K/V in the MHA path, and absorbs the same matrices in the MQA path. +Float64 makes the equality check sensitive enough to catch a transposed index or +an incorrect contraction. + +```python +import math +import torch + +torch.manual_seed(0) +Q, K, H = 3, 5, 4 +D_CONTENT, D_LATENT, D_VALUE, D_ROPE, D_MODEL = 7, 6, 8, 3, 11 + +# Per-head queries, one shared latent KV per key token, and shared RoPE keys. +q_content = torch.randn(Q, H, D_CONTENT, dtype=torch.float64) +q_rope = torch.randn(Q, H, D_ROPE, dtype=torch.float64) +c_kv = torch.randn(K, D_LATENT, dtype=torch.float64) +k_rope = torch.randn(K, D_ROPE, dtype=torch.float64) + +W_UK = torch.randn(H, D_CONTENT, D_LATENT, dtype=torch.float64) +W_UV = torch.randn(H, D_VALUE, D_LATENT, dtype=torch.float64) +W_O = torch.randn(D_MODEL, H, D_VALUE, dtype=torch.float64) + +# MHA mode: explicitly expand a key and value for every head. +k_content = torch.einsum("hdc,kc->khd", W_UK, c_kv) +v_content = torch.einsum("hvc,kc->khv", W_UV, c_kv) +scale = 1.0 / math.sqrt(D_CONTENT + D_ROPE) +scores_mha = ( + torch.einsum("qhd,khd->qhk", q_content, k_content) + + torch.einsum("qhr,kr->qhk", q_rope, k_rope) +) * scale +prob = torch.softmax(scores_mha, dim=-1) +head_out_mha = torch.einsum("qhk,khv->qhv", prob, v_content) +model_out_mha = torch.einsum("mhv,qhv->qm", W_O, head_out_mha) + +# MQA mode: move W_UK to Q, attend to c_kv, then move W_UV to output. +q_absorbed = torch.einsum("qhd,hdc->qhc", q_content, W_UK) +scores_mqa = ( + torch.einsum("qhc,kc->qhk", q_absorbed, c_kv) + + torch.einsum("qhr,kr->qhk", q_rope, k_rope) +) * scale +latent_out = torch.einsum("qhk,kc->qhc", torch.softmax(scores_mqa, -1), c_kv) +W_O_absorbed = torch.einsum("mhv,hvc->mhc", W_O, W_UV) +model_out_mqa = torch.einsum("mhc,qhc->qm", W_O_absorbed, latent_out) + +torch.testing.assert_close(scores_mha, scores_mqa, rtol=1e-12, atol=1e-12) +torch.testing.assert_close(model_out_mha, model_out_mqa, rtol=1e-12, atol=1e-12) +print("weight absorption: exact up to float64 roundoff") +``` + +Notice that the scale remains the scale of the model's semantic QK head, not +$1/\sqrt{D_{latent}}$ merely because the absorbed dot product happens to have +$D_{latent}$ coordinates. + +:::{admonition} The query side can also be compressed +:class: note + +The query projection may use a separate low-rank latent as well. Write its +down- and up-projections as $D_Q$ and $U_Q$: + +$$ +c^Q=D_Qh,\qquad q^C=U_Qc^Q. +$$ + +This factorization mainly reduces activation memory during training; it does +not shrink the KV cache further. It happens before core attention: the `q` +tensor presented to the attention implementation is already projected, so the +KV-path analysis treats $q^C$ as an input. +::: + +For the complete MLA architecture, training design, and original notation, see +the [DeepSeek-V2 paper](https://arxiv.org/abs/2405.04434). The discussion here +keeps the KV compression, decoupled RoPE, and weight absorption needed to +understand the FlashMLA kernel that follows. + +Weight absorption explains how one shared KV row can serve many query heads. +Sparse prefill introduces a separate boundary between token selection and the +sparse core-attention operator that consumes the selected rows. + +## Token-selection boundary of the sparse-prefill operator + +Dense attention visits every eligible KV token. DSA first uses a lightweight +*lightning indexer* to score candidate tokens and chooses a top-$k$ set for each +query. The sparse core attention then reads only those latent KV entries. If the +original context length is $L$, this changes the core-attention work from +$O(L^2)$ to $O(Lk)$ for prefill, although the indexer has its own cost. + +```{figure} ../img/flashmla_sparse_story.png +:width: 100% +:alt: A lightning indexer selecting token rows before the sparse-prefill attention operator + +Selection and attention are separate operators. The indexer produces row +addresses; the sparse-prefill operator gathers those rows and performs the +QK--softmax--PV computation. +``` + +The sparse-prefill operator receives the indexer's result as an `indices` +tensor. Its semantic contract is: + +1. gather the requested KV rows; +2. mark out-of-range and length-masked positions invalid; +3. compute attention over the remaining rows, including duplicates if the caller + supplied duplicates; +4. return the output, maximum logit, and log-sum-exp. + +There is no causal flag in this interface. A caller that requires causal +attention must produce an index list containing only allowed keys. Sparsity is +not itself a causal mask. + +When `topk_length` is present, every query must satisfy +`0 <= topk_length[q] <= topk`. This is a caller precondition. + +This prefill interface also has no batch dimension. Each query token supplies +one selected-token list, and all `h_q` query heads share that list. A serving +system must flatten or otherwise map batches before making this call. + +An executable CPU reference makes these rules testable independently of any GPU +implementation. + +## An executable sparse-attention reference + +The general shape notation is: + +| Symbol | Meaning | +| --- | --- | +| `s_q` | number of query rows | +| `s_kv` | number of addressable KV rows | +| `h_q` | query heads in each query row | +| `h_kv` | KV heads in each KV row; this interface requires 1 | +| `d_qk` | query/key width used by QK | +| `d_v` | value and output width, with `d_v <= d_qk` | +| `topk` | index slots supplied for each query row | + +The corresponding tensors are `q[s_q,h_q,d_qk]`, +`kv[s_kv,h_kv,d_qk]`, `indices[s_q,1,topk]`, and +`out[s_q,h_q,d_v]`. The optional sink has shape `[h_q]`, the optional +`topk_length` has shape `[s_q]`, and both returned statistics have shape +`[s_q,h_q]`. The general sparse-prefill contract already requires `h_kv=1`; +the regular head-128 specialization additionally fixes `h_q=128` and +`d_v=512`. + +For the absorbed MQA contract, `kv[:, 0, :]` supplies both K and V: all +`d_qk` coordinates participate in QK, while the first `d_v` coordinates are +the latent value. Here `sm_scale` is the model's semantic QK scale. + +Invalid indices must be clamped *before* a PyTorch gather and then masked out. +Directly indexing with `-1` would incorrectly select the last row. + +The gathered V row for an out-of-range address is also cleared before PV, +because a zero softmax weight does not neutralize a NaN under IEEE arithmetic. + +An attention sink is equivalent to adding a logit whose value vector is zero. +Fix one query and head. Let $x_j$ be the ordinary scaled logit for selected KV +row $j$, $v_j$ its value, $a$ the sink logit, and $m$ the maximum of the +ordinary logits. The sink changes only the output denominator: + +$$ +O=\frac{\sum_j e^{x_j-m}v_j} +{\sum_j e^{x_j-m}+e^{a-m}}. +$$ + +The sink has no effect on `max_logits` or the returned log-sum-exp (`lse`). If every selected +position is invalid, the operator instead uses the explicit convention +`out=0`, `max_logits=-inf`, and `lse=+inf`. + +The executable oracle follows four stages: + +1. validate and safely gather the requested rows; +2. construct the validity predicate and clear only out-of-range V sentinels; +3. compute scaled logits, unnormalized weights, numerator, denominator, and the + optional sink term; and +4. normalize the output and return the two statistics, including the all-invalid + convention. + +The CPU implementation is: + +```python +import math +import torch + +def sparse_prefill_reference( + q, kv, indices, sm_scale, d_v, *, attn_sink=None, topk_length=None +): + """Reference for q[SQ,H,D], kv[SKV,1,D], indices[SQ,1,TOPK].""" + s_q, h_q, d_qk = q.shape + s_kv, h_kv, kv_width = kv.shape + assert h_kv == 1 and kv_width == d_qk and d_v <= d_qk + assert indices.shape[:2] == (s_q, 1) + + idx = indices[:, 0].to(torch.long) # [SQ, TOPK] + topk = idx.shape[1] + assert topk > 0 + in_range = (idx >= 0) & (idx < s_kv) + safe_idx = idx.clamp(0, s_kv - 1) + focused_kv = kv[:, 0].float()[safe_idx] # [SQ, TOPK, D] + + # OOB indices use a clamped boundary row only as a safe gather address. Clear + # their V rows before PV so that 0 * NaN cannot leak from that sentinel row. + # Only OOB sentinel rows need clearing for the defined-input oracle. + focused_v = torch.where( + in_range[:, :, None], + focused_kv[:, :, :d_v], + torch.zeros_like(focused_kv[:, :, :d_v]), + ) + + position = torch.arange(topk, device=q.device)[None, :] + if topk_length is not None: + assert topk_length.shape == (s_q,) + assert bool(((0 <= topk_length) & (topk_length <= topk)).all()) + length = ( + topk_length.to(torch.long)[:, None] + if topk_length is not None + else torch.full((s_q, 1), topk, device=q.device) + ) + valid = in_range & (position < length) + + logits = torch.einsum("qhd,qkd->qhk", q.float(), focused_kv) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], -torch.inf) + max_logits = logits.amax(dim=-1) + have_valid = valid.any(dim=-1)[:, None] + + # Avoid (-inf)-(-inf) on rows for which every selected index is invalid. + softmax_origin = torch.where(have_valid, max_logits, torch.zeros_like(max_logits)) + weight = torch.exp(logits - softmax_origin[:, :, None]) + weight = torch.where(valid[:, None, :], weight, torch.zeros_like(weight)) + denominator = weight.sum(dim=-1) + numerator = torch.einsum("qhk,qkv->qhv", weight, focused_v) + + if attn_sink is None: + sink_term = torch.zeros_like(denominator) + else: + sink_term = torch.exp(attn_sink.float()[None, :] - softmax_origin) + out = numerator / (denominator + sink_term).clamp_min(torch.finfo(torch.float32).tiny)[ + :, :, None + ] + out = torch.where(have_valid[:, :, None], out, torch.zeros_like(out)) + + # FlashMLA's reported LSE excludes the attention sink. Its all-invalid + # convention is max_logits=-inf, lse=+inf, output=0. + lse = torch.where( + have_valid, + max_logits + torch.log(denominator), + torch.full_like(max_logits, torch.inf), + ) + return out, max_logits, lse + + +torch.manual_seed(1) +q = torch.randn(2, 3, 6) +kv = torch.randn(9, 1, 6) +indices = torch.tensor([[[0, 3, -1, 12]], [[8, 1, 4, 2]]], dtype=torch.int32) +topk_length = torch.tensor([3, 2], dtype=torch.int32) +attn_sink = torch.randn(3) +out, max_logits, lse = sparse_prefill_reference( + q, kv, indices, 1 / math.sqrt(6), 4, + attn_sink=attn_sink, + topk_length=topk_length, +) +assert out.shape == (2, 3, 4) +assert max_logits.shape == lse.shape == (2, 3) +assert torch.isfinite(out).all() + +# An OOB sentinel must not inherit NaNs from the row used as its safe address. +nan_q = torch.ones(1, 1, 2) +nan_kv = torch.tensor([[[torch.nan, torch.nan]], [[2.0, 3.0]]]) +nan_indices = torch.tensor([[[-1, 1]]], dtype=torch.int32) +nan_out, _, _ = sparse_prefill_reference(nan_q, nan_kv, nan_indices, 1.0, 2) +torch.testing.assert_close(nan_out, torch.tensor([[[2.0, 3.0]]])) +print(out.shape, max_logits.shape, lse.shape) +``` + +The reference makes the sparse-prefill operator's numerical contract executable +for defined inputs and return conventions. An implementation supplies tile +partitioning, storage reuse, and pipeline overlap while preserving that +contract. + +## Sparse prefill in the FlashMLA operator family + +FlashMLA spans both the sequence stage and the selection pattern. *Decoding* +adds a new query (or a small speculative group) step by step while reusing the +KV cache. The +[official FlashMLA repository](https://github.com/deepseek-ai/FlashMLA) +organizes its operators and their implementations into four broad families: + +| Selection | Sequence stage | Representative purpose | +| --- | --- | --- | +| dense | prefill | MHA forward and backward | +| dense | decoding | read an MLA KV cache for newly generated queries | +| token-sparse | prefill | DSA core attention over a selected token list | +| token-sparse | decoding | DSA inference over a selected FP8 KV cache | + +The token-sparse prefill cell combines externally selected, irregular KV-row +addresses with core attention, creating the Blackwell scheduling problem +studied here. Sparse decode uses a different paged-cache, scheduling, and +reduction contract. + +The FlashMLA sparse-prefill interface is conceptually + +```text +out, max_logits, lse = flash_mla_sparse_fwd( + q, kv, indices, sm_scale, + d_v=512, + attn_sink=attn_sink, # optional [h_q], float32 + topk_length=topk_length, # optional [s_q], int32 +) +``` + +The tensor dimensions and explicit arguments have the semantics exercised by +the oracle: `h_q` is the query-head count, `s_q` is the query-row count, +`sm_scale` scales QK scores, `d_v` selects the value/output width, `attn_sink` +optionally adds one zero-valued logit per query head, and `topk_length` limits +each query's valid index prefix. The call returns the normalized output, maximum +scaled logit, and log-sum-exp without the sink. + +The TIRx implementation exposes a registry/dispatch bridge with the same +`flash_mla_sparse_fwd` name; it is not a complete replica of the FlashMLA Python +API above. By shape, this bridge selects one of three SM100 phase-1 +specializations of the same sparse-prefill computation. + +:::{admonition} The TIRx entry specializes `sm_scale` +:class: warning + +The FlashMLA interface above accepts `sm_scale` at runtime. The TIRx dispatch +entry exposes a narrower call signature: all three prefill specializations bind +the scale to `1 / sqrt(d_qk)`, and the launch application binary interface +(ABI) has no scale argument. + +Passing `sm_scale=...` through the `**kwargs` wrappers is silently ignored. The +B200 examples in this chapter therefore validate the computation only at +`sm_scale = 1 / sqrt(d_qk)`; they do not demonstrate runtime-scale parity with +the complete FlashMLA interface. + +A model whose semantic QK scale differs must expose or specialize the correct +value. Weight absorption does not justify changing that scale. +::: + +Two more boundary conditions belong specifically to this TIRx path. Its +prefill specializations use `topk_length` to decide how many selected-row tiles +to visit but do not clip it, so a value above `topk` can step beyond the logical +`indices` storage. They also do not promise to sanitize NaNs in an in-range KV +row that lies beyond `topk_length`; callers should keep such length-masked rows +finite. Within these preconditions and at the specialized scale, every +dispatched implementation must reproduce the reference contract. + +## The Blackwell regular head-128 case + +The regular head-128 case retains the QK--softmax--PV chain and adds irregular +gather, absorbed latent KV, and cooperative thread-block ownership. + +### Shape and dispatch conditions for regular head-128 + +The regular head-128 module has this shape-specialized signature: + +| Tensor | Shape | Type | Meaning | +| --- | --- | --- | --- | +| `q` | `[s_q, 128, d_qk]` | BF16 | absorbed queries | +| `kv` | `[s_kv, 1, d_qk]` | BF16 | shared latent/positional KV rows | +| `indices` | `[s_q, 1, topk]` | int32 | direct KV row indices | +| `attn_sink` | `[128]` | FP32 | optional per-head sink logits | +| `topk_length` | `[s_q]` | int32 | optional valid prefix length; every entry is in `[0, topk]` | +| `out` | `[s_q, 128, 512]` | BF16 | sparse-attention result | +| `max_logits` | `[s_q, 128]` | FP32 | maximum scaled logit | +| `lse` | `[s_q, 128]` | FP32 | natural-log sum-exp, without sink | + +In this specialization, 128 is the query-head count, while the 1 in `kv` means +that every head shares the same KV row. A common $d_{qk}=576$ case combines 512 +latent-content coordinates with 64 RoPE coordinates; `d_v=512` is the +latent-value width. Other MLA operators may use different shapes. + +For this MLA layer, the equivalent MHA representation has QK feature width +$128+64=192$ and value/output width 128. The absorbed MQA representation used +here has widths $512+64=576$ and 512. A rough count of multiply-add +coordinates per query--key pair is therefore $192+128=320$ versus +$576+512=1088$, about 3.4 times as many for the absorbed representation. +This is only an arithmetic-width intuition, not a prediction of kernel runtime. + +The running shape is `s_q=1`, +`s_kv=8192`, `h_q=128`, `h_kv=1`, `d_qk=576`, `d_v=512`, and `topk=2048`. It is +one query row with 128 query heads sharing 2048 selected-index slots. Those +slots may contain duplicate or out-of-range addresses. Without a shorter +`topk_length`, the physical schedule visits $N=16$ tiles of 128 slots. When +`topk_length` is present, it visits +`max(ceil(topk_length / 128), 1)` tiles. + +Each selected-index tile follows six semantic steps. In the source notation, +$L$ denotes raw QK logits, $W$ denotes BF16 unnormalized exponential weights, +`mi` is the online-softmax exponent origin, `li` is the denominator accumulated +relative to that origin, and $\widetilde O$ is the accumulated output before +division by the denominator: + +```text +for each 128-slot selected-index tile: + 1. gather the tile's K rows + 2. gather its V rows and build the validity mask + 3. QK: 128 query heads x 128 selected-index slots -> logits L + 4. update masked online softmax, producing W, origin mi, and denominator li + 5. rescale the running state when needed, then accumulate O~ += W @ V +after all selected-index tiles: + 6. normalize O~ by li + sink; store out, max_logits, and lse +``` + +In the running example, the first five steps repeat 16 times before the sixth +step runs. + +### Complete source navigation + +The chapter does not reproduce the entire device function. Instead, it presents +short excerpts organized around QK, softmax, PV, data movement, and +synchronization. Read the complete source through these entry points: + +| Goal | Source entry point | +| --- | --- | +| Unified entry and shape dispatch | [`flash_mla_sparse_fwd.py` lines 66--125](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/flash_mla_sparse_fwd.py#L66-L125) | +| Config, test data, PyTorch reference, and launch ABI | [`sparse_prefill_head128_phase1.py` lines 66--244](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L66-L244) | +| Complete regular head-128 device kernel | [`sparse_prefill_head128_phase1.py` lines 247--865](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L247-L865) | +| CTA-pair TMA, tcgen05 MMA, and validity-mask helpers | [`_tma.py` lines 10--60](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/_tma.py#L10-L60), [`_gemm.py` lines 8--29](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/_gemm.py#L8-L29), and [`_mask.py` lines 10--29](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/_mask.py#L10-L29) | +| Specialization, compilation, launch, and numerical checks | [`sparse_prefill_head128_phase1.py` lines 868--905](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L868-L905) | + +A productive order is dispatch and tensor ABI first, then the WG0, WG1, WG2, +and WG3 branches in `_kernel`. Follow TMA, MMA, and mask calls into their helpers +only when they appear, and finish with `run_test` to connect inputs, outputs, and +the reference. Each excerpt in the chapter preserves the source variable names +and slices and links to its full context. + +The TIRx regular head-128 implementation is in +[`sparse_prefill_head128_phase1.py`](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py). +TIRx extends TVM 0.26's TIR Python DSL: `T` denotes the TIR script namespace, +while `Tx` contains GPU-kernel helpers. The `phase1` name follows the +corresponding CUDA implementation's file naming; on this regular prefill path, +one kernel produces the complete `(out, max_logits, lse)` result. + +The implementation uses three execution levels. A cooperative thread array +(CTA) is one CUDA thread block. Two adjacent CTAs form a cluster that can +participate in CTA-group tensor-core operations. A warpgroup is four warps, or +128 threads, assigned one specialized role. Here one two-CTA cluster owns one +query row. + +In the mathematical introduction, $p$ denoted a normalized softmax +probability. In the source, however, `tmem_p` and the register variable `p` +hold the raw QK logits already named $L$. The source's `s_frag` and +`s_smem_gemm` hold the unnormalized exponential weights already named $W$. +Only the epilogue divides the accumulated output by `li` plus the optional sink +term. + +Short TIRx code blocks are contextual excerpts from the linked implementation. +Independently runnable blocks are labeled explicitly. + +The following constants describe the tile sizes, thread count, and +synchronization slots: + +```python +B_H = 128 +B_TOPK = 128 +D_V = 512 +NUM_BUFS = 2 +NUM_THREADS = 512 +D_TQ = 384 +``` + +`B_H` is the 128-head logical tile, `B_TOPK` is the 128 selected-index slots +processed per streaming tile, `D_V` is the value/output width, and `D_TQ` is the +384-coordinate Q suffix moved to dedicated on-chip storage. `NUM_THREADS=512` +gives each CTA four warpgroups. `NUM_BUFS=2` provides two slots for +synchronization state and two small packed-validity-mask slots. + +The regular head-128 specialization accepts `d_qk` 512 or 576, requires +`h_kv=1`, `d_v=512`, and requires `topk` to be a positive multiple of 128. The +TIRx dispatch entry makes one additional choice for 128 heads: `d_qk=512` +with `topk<=1280` selects the small-top-k head-128 specialization; other +supported head-128 shapes select this regular head-128 specialization. Head-64 +shapes select the head-64 specialization. + +The positive-`topk` condition is a required caller precondition. The TIRx +dispatch entry rejects `topk<=0`, but the per-specialization `_cfg().validate()` +methods check divisibility without checking positivity. A direct import of one +specialization must therefore still reject or avoid nonpositive `topk`; +acceptance by that local validator does not make such a launch valid. + +Once `tirx-kernels` is installed, the dispatch can be inspected without +launching a GPU kernel: + +```python +from tirx_kernels.flashmla.flash_mla_sparse_fwd import ( + dispatch_reason, + select_kernel, +) + +shape = dict(h_q=128, h_kv=1, d_qk=576, d_v=512, topk=2048) +assert select_kernel(**shape) == "sparse_flashmla_prefill_head128_phase1" +print(dispatch_reason(**shape)) +# sm100 h_q=128 dispatches to regular head128 phase1 +``` + +The dispatch itself is documented in +[`flash_mla_sparse_fwd.py`](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/flash_mla_sparse_fwd.py#L66-L120). +The registry keeps dispatch separate from the device schedule, preserving the +distinction between operator selection and schedule implementation. + +The selected implementation uses two thread blocks to cooperate on the six +steps without duplicating the whole tile. + +### Query-row ownership across two CTAs + +The difficulty lies in steps 3 and 5 of the tile skeleton. QK must form every +pair of 128 query heads and 128 selected tokens; PV then contracts the token +axis to produce 512 value coordinates. The implementation therefore lets two +CTAs form one logical tile and changes their partition axis between QK and PV. +The ownership map is: + +```{figure} ../img/flashmla_cta_ownership.png +:width: 100% +:alt: Two-CTA ownership of query heads, selected K rows, and V feature columns + +For each query row, the CTA pair changes its logical partition between QK and +PV. Each CTA ultimately writes 64 complete output-head vectors, each with 512 +coordinates. +``` + +The pair divides three different axes in three different ways: + +| Resource within a 128-token top-k tile | CTA 0 | CTA 1 | +| --- | --- | --- | +| query/output head ownership | heads 0--63 | heads 64--127 | +| K-row gather ownership | selected tokens 0--63 | selected tokens 64--127 | +| V-feature gather ownership | value columns 0--255 | value columns 256--511 | + +A 2-CTA tensor-core operation forms one collective logical tile. QK computes the +cross-product of 128 heads and 128 selected tokens; PV then contracts those +tokens into 512 value coordinates. The partition rotates between the two GEMMs, +supported by collective `cta_group=2` MMA, paired on-chip layouts, and cross-CTA +synchronization. + +The launch topology implements this ownership map. The grid contains +`2 * s_q` CTAs and clusters adjacent CTAs in pairs: + +```python +block_idx = T.cta_id([2 * s_q]) +T.cta_id_in_cluster([2]) +cta_idx: T.let = block_idx % 2 +s_q_idx: T.let = block_idx // 2 +thread_idx = T.thread_id([512]) +T.warpgroup_id([4]) +``` + +One cluster therefore owns one query row, and each CTA has four warpgroups. The +same division appears in the data indexing: Q is chunked by `cta_idx`, the K +producer selects the `cta_idx` half of every top-k block, and the V producer +starts at `cta_idx * 256`. + +## Tile residency and lifetime + +The ownership map says *who* computes each piece; residency says where a piece +waits between producers and consumers: + +- global memory (GMEM) holds the input and output tensors; +- shared memory (SMEM) is the ordinary on-chip scratchpad visible to a CTA; +- Blackwell tensor memory (TMEM) is a separate on-chip space near the tensor + cores, used for operands and large accumulators. + +```{figure} ../img/flashmla_dataflow.png +:width: 100% +:alt: Data residency and lifetime reuse across global memory, shared memory, tensor memory, and WG0 registers during QK, softmax, PV, and the epilogue + +Q is split between an SMEM prefix and a TMEM suffix. Gathered K/V enter SMEM, +raw QK logits and the output accumulate in TMEM, and unnormalized softmax +weights cross back through SMEM for PV. +``` + +The arrows in the figure correspond to two Blackwell mechanisms. The Tensor +Memory Accelerator (TMA) moves data asynchronously between GMEM and SMEM and +provides the sparse `gather4` path used here. The `tcgen05` tensor-core +instruction family reads operands from SMEM or TMEM and keeps its large +accumulators in TMEM. TMEM and SMEM have complementary roles: TMA gathers land +in SMEM, and the softmax warpgroup materializes BF16 unnormalized weights there +for PV. + +The source abbreviates tensor-core operand residency as **SS** when both matrix +operands come from SMEM and **TS** when the first comes from TMEM and the second +from SMEM. The Q prefix follows the SS path; its suffix follows the TS path. + +For one CTA, the important logical views are: + +| Storage | Logical tile | Lifetime and purpose | +| --- | --- | --- | +| SMEM `q_full` | `64 x d_qk` BF16 | Q prologue; its prefix remains for SS QK | +| TMEM `q_tmem` | `64 x 384` BF16 | suffix of Q used by TS QK | +| SMEM `k_smem` | `64 x d_qk` BF16 | this CTA's gathered half of a 128-row K tile | +| TMEM `tmem_p` | `64 x 128` FP32 logical view | raw QK logits $L$ consumed by softmax | +| SMEM `s_smem_gemm` | `64 x 128` BF16 | unnormalized exponential weights $W$ for PV | +| SMEM `v_smem_gemm` | `128 x 256` BF16 logical view | rearranged view of `v_smem`: all tile rows, this CTA's V columns | +| TMEM `o_tmem` | `64 x 512` FP32 logical view | running unnormalized output | +| SMEM `o_smem` | `64 x 512` BF16 | epilogue staging before TMA store | + +These are logical views; CTA-group TMEM layouts and rearrangements give the MMA +and load/store instructions their physical lane mapping. The source allocates +one 512-column CTA-group TMEM pool, then carves out O, raw-logit, and Q views. + +The source object named `SMEMPool` describes the corresponding shared-memory +allocation. SMEM is aggressively aliased: `q_full`, the gathered K/V region, +and the output epilogue reuse a union-like base when their lifetimes permit. + +After the final 384 Q columns have moved to TMEM, only the +$d_{sq}=d_{qk}-384$ prefix must stay live for the first QK part. For +`d_qk=512`, $d_{sq}=128$; for 576, it is 192. + +The allocation plan is visible in +[`sparse_prefill_head128_phase1.py` lines 302--365](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L302-L365). + +Safe in-place reuse of these storage regions relies on a **completion barrier**, +a small hardware state object that records when an asynchronous producer has +finished. Its phase bit distinguishes successive uses of each synchronization +slot, allowing released portions of the single K/V/$L$/$W$ tile storage to be +overwritten safely. + +## Warpgroup responsibilities + +The four warpgroups have specialized roles: + +| Warpgroup | Warps | Responsibility | +| --- | --- | --- | +| WG0 | 0--3 | load raw logits $L$ from TMEM, mask, online softmax, write weights $W$, rescale O, epilogue | +| WG1 | 4--7 | load index fragments and issue gather4 TMA for K | +| WG2 | 8--11 | load index fragments and issue gather4 TMA for V | +| WG3 | 12--15 | warp 12 of CTA 0 issues CTA-group QK/PV MMA; warp 13 in each CTA builds validity masks | + +WG3 concentrates its active responsibilities in warp 12 for asynchronous MMA +issue and warp 13 for validity masks. This asymmetric role assignment matches +the parallelism of each operation: one elected lane issues MMA for the CTA pair, +warp 13 packs validity bits, and WG0 uses many lanes for exponentiation, row +reductions, and epilogue conversion. + +:::{admonition} Role-specific register limits +:class: note + +The register budgets match the roles. WG0 raises its limit to 144 registers and +WG3 to 168; producer groups lower theirs to 96. The TIRx API spells these calls +as `T.ptx.setmaxnreg(True, ...)` for an increase and +`T.ptx.setmaxnreg(False, ...)` for a decrease. +::: + +The gather handoff follows a ready/free protocol. A producer waits until a +tile's storage is **free**, writes or asynchronously fills the tile, and signals +**ready**. The consumer waits for ready, uses the tile, and signals free when the +storage may be overwritten. These barriers transfer ownership of the in-place +storage. + +### From irregular rows to regular tiles + +Sparse row addresses destroy the contiguous 2-D copy pattern used by dense +attention. WG1 and WG2 use explicit TMA `gather4`: one issue supplies exactly +four row coordinates, so a warp can bring noncontiguous KV rows into a regular +SMEM tile. + +The contextual excerpt identifies the addresses read by one `gather4` issue and +the barrier that receives its completion. Its index names and slices match the +linked implementation: + +- `gather4=[...]` supplies the four KV source-row coordinates for this issue; +- `cur_buf = k % NUM_BUFS` selects the current slot in the two-slot barrier ring; + and +- `bar` is the ready-barrier array passed to this copy helper. `leader_mbar` + selects the CTA-pair leader's slot, where TMA reports asynchronous completion. + +The remaining names describe the surrounding layouts. + +```python +_kv_gather_tma = partial( + tma_config, + dispatch="tma_explicit", + cta_group=2, + cta_mask=T.uint16(1), + cache_hint=T.uint64(0x14F0000000000000), +) + +for row_group in T.unroll(WG1_ROWS_PER_WARP): + for col_atom in T.unroll(col_count): + col = T.meta_var((col_start + col_atom) * 64) + Tx.copy_async( + k_gather_tile[ + row_group * 4 : row_group * 4 + 4, + col_atom * 64 : col_atom * 64 + 64, + ], + kv_tma[0:1, col : col + 64], + **_kv_gather_tma( + mbar=leader_mbar(bar.ptr_to([cur_buf])), + gather4=[indices_int4[row_group, lane] for lane in range(4)], + ), + ) +``` + +Gathering and validity are related but separate. Warp 13 loads eight indices per +active lane and calls `pack_valid_mask8`. Bit $i$ is one only if + +$$ +0\leq\text{index}_i 0: + sq_smem = q_full.sub[:, :d_sq] + Tx.gemm_async( + tmem_p[:, :], + sq_smem[:, :d_sq], + k_smem[:, :d_sq], + **_mma_config(accum=mma_p_accumulate, smem_desc=mma_smem_desc), + ) + mma_p_accumulate = T.uint32(1) + +Tx.gemm_async( + tmem_p[:, :], + q_tmem[:, :D_TQ], + k_smem[:, d_sq : d_sq + D_TQ], + **_mma_config(accum=mma_p_accumulate, smem_desc=mma_smem_desc), +) +``` + +The first operation is SS: both Q and K operands are described from SMEM. The +second is TS: Q's 384-column suffix comes from TMEM while K remains in SMEM. Both +write the same FP32 raw-logit accumulator (`tmem_p` in the source); the first +clears it and the second accumulates. Splitting Q this way preserves only the +smaller Q prefix in SMEM, making the union allocation possible without giving +up the TS path for the larger suffix. + +After softmax, PV is an SS GEMM: BF16 $W$ and V are both in SMEM, while the FP32 O +accumulator stays in TMEM. The kernel splits the V rows and output columns into +two halves each; the four combinations collectively update all 512 value +coordinates. + +## Lazy O rescaling in online softmax + +QK has produced raw logits $L$; softmax must now turn each tile into $W$ while +preserving state from earlier selected tiles. In the running shape, this is why +the state must be merged across $N=16$ tiles. For every raw QK dot product $x_j$ +in the next tile, the kernel first computes + +$$ +r_j=x_j\cdot\text{semantic\_QK\_scale}\cdot\log_2(e). +$$ + +The TIRx specialization binds the multiplier named `sm_scale_div_log2` to +`(1 / sqrt(d_qk)) * log2(e)`. Despite that source name, each $r_j$ is simply a +model-scaled score expressed in base-2 units so that the kernel can use `exp2`. + +For a stream of such score tiles, online softmax stores a base-2 row origin $m$, +denominator $\ell$, and unnormalized output $\widetilde O$. To merge the next +tile, + +$$ +m'=\max(m,\max_j r_j),\qquad +\alpha=2^{m-m'}, +$$ + +$$ +\ell'=\alpha\ell+\sum_j2^{r_j-m'},\qquad +\widetilde O'=\alpha\widetilde O+\sum_j2^{r_j-m'}v_j. +$$ + +Rescaling the full 512-coordinate O tile whenever the row maximum increases +would be expensive. The lazy-threshold excerpt maps to the recurrence as +follows: + +- `cur_pi_max` is the current tile's maximum in the base-2 exponent domain; +- `mi` is the retained numerical origin $m$; +- `real_mi` is the exact maximum retained for the reported `max_logits`; +- `li` is the running denominator $\ell$; and +- `attn_sink_log2` is the optional sink logit in the same base-2 domain. + +The head-128 kernel then makes one warp-uniform decision: + +```python +should_scale_o: T.bool = ( + T.ptx.any_sync(T.uint32(0xFFFFFFFF), cur_pi_max - mi > 6.0) != 0 +) + +if not should_scale_o: + scale_for_old = 1.0 + new_max = mi +else: + new_max = T.max(cur_pi_max, mi) + scale_for_old = T.ptx.exp2(mi - new_max) +``` + +If the new tile maximum is at most 6 base-2 units above the stored origin, the +kernel keeps the old origin. New exponentials may then be as large as $2^6=64$, +but the accumulated O does not need a rescale. Once the difference exceeds 6, +it rebases and rescales both $\ell$ and, when it already exists, O. The warp-wide +`any_sync` keeps the decision uniform for the participating rows. + +Because `real_mi` is maintained separately, the optimization does not change +`max_logits`. At the end, the two 64-token contributions to each logical row are +combined, and the kernel emits + +$$ +\mathrm{lse}=m\ln 2+\ln\ell. +$$ + +The optional attention sink changes the final output scale to + +```python +output_scale: T.float32 = T.cuda.fdividef( + T.float32(1.0), li + T.ptx.exp2(attn_sink_log2 - mi) +) +``` + +but deliberately leaves the reported LSE untouched. All-invalid rows are +special-cased to output zero with `max_logits=-inf` and `lse=+inf`, matching the +executable reference. + +## Race-free pipeline overlap + +The pipeline repeats the ready/free ownership handoffs across tiles. QK must +finish before softmax consumes its logits, and PV must wait until softmax has +produced its weights. Meanwhile, the gather warpgroups move forward whenever an +in-place K or V segment becomes reusable. + +```{figure} ../img/flashmla_pipeline_stages.png +:width: 100% +:alt: Sparse-prefill pipeline fill, steady state, and drain, including the overlap among QK, softmax, and PV for adjacent tiles + +The fill iteration issues QK(0) without a previous PV, while the drain iteration +issues PV($N-1$) without a new QK before the final epilogue. In steady state, +softmax($k-1$) may overlap QK($k$). The sole MMA issuer then issues PV($k-1$) +after QK($k$); once QK($k$) completes, softmax($k$) may overlap that PV. +Producer warpgroups gather the next safe K/V segments around this serial issuer +order. For the running shape, $N=16$. +``` + +At the coarsest level, four ownership handoffs repeat for each tile: + +1. WG1/WG2 publish gathered K/V segments, while warp 13 publishes validity; +2. WG3 completes QK, publishes $L$ to WG0, and returns consumed K segments; +3. WG0 masks $L$, updates online softmax, publishes $W$ to WG3, and releases the + raw-logit and validity storage; and +4. WG3 completes PV, returning V storage to WG2 and making O safe for WG0 to + rescale or finalize. + +An **mbarrier** is the hardware completion object behind these handoffs. It +tracks expected arrivals or TMA bytes and carries a phase, so a wait identifies +the intended reuse of a slot. The four-step view gives the causal chain; the +part-level edges release K and V segments as early as their consumers finish. + +The kernel initializes its mbarriers in warp 0, performs a cluster sync, launches +the Q prologue, allocates CTA-group TMEM, and then enters specialized loops. + +TMA completion for CTA-group gathers is routed to a named leader barrier, so +issues from the pair contribute to one expected byte count. + +The main barrier edges are easier to read as ownership transfers: + +| Barrier | Producer to consumer | Storage protected | +| --- | --- | --- | +| `bar_k_part0_ready` | WG1 to WG3 | K prefix for SS QK | +| `bar_qk_part_done` | WG3 to WG1 | permission to overwrite K prefix after SS QK completion | +| `bar_k_part1_ready` | WG1 to WG3 | K suffix for TS QK | +| `bar_qk_done` | WG3 to WG0 and WG1 | raw logits $L$ ready; K suffix reusable after QK completion | +| `bar_p_free` | WG0 to WG3 | TMEM raw-logit tile consumed before next overwrite | +| `bar_k_valid_ready/free` | warp 13 to/from WG0 | packed validity mask | +| `bar_so_ready` | WG0 to WG3 | BF16 weights $W$ ready for PV | +| `bar_v_part0_ready` / `bar_sv_part_done` | WG2 to/from WG3 | first V half | +| `bar_v_part1_ready` / `bar_sv_done` | WG2 to WG3, then WG3 to WG2 and WG0 | second V half; PV/O completion before V reuse, O rescale, or epilogue | + +The ring index is + +```python +cur_buf = k % 2 +cur_phase = (k // 2) & 1 +``` + +so a reused barrier slot can distinguish a new arrival from one made two +iterations earlier. + +`bar_qk_part_done` allows the producer to replace K's prefix before its suffix +is reusable. The two `bar_sv_*` edges do the analogous job for V. + +```{figure} ../img/flashmla_pipeline.png +:width: 100% +:alt: Detailed sparse-prefill pipeline showing serial QK and PV issue, part-wise K and V reuse, the mask-slot ring, and the WG0 handoff + +This detailed view names the part-level reuse edges and the mask-slot ring. +Barrier phases protect reuse of the single in-place tile storage. +``` + +The memory model has one more distinction. Ordinary thread loads and stores see +SMEM through the **generic proxy**; TMA and tensor-core asynchronous accesses use +an **asynchronous proxy**. A barrier reports completion, but completion alone +does not establish the required visibility and ordering between those proxies. + +Two kinds of fences therefore appear around the barrier edges. +`T.ptx.tcgen05.fence.*` orders TMEM accesses relative to thread-visible work, +while `T.ptx.fence.proxy_async("shared::cta")` establishes cross-proxy ordering +between generic and asynchronous accesses to SMEM. + +Here the proxy fence is needed both when generic stores of $W$ or the epilogue +tile precede tcgen05/TMA asynchronous reads, and after an asynchronous SMEM read +completes before generic code overwrites aliased storage. + +Thus an mbarrier communicates completion and an ownership handoff, whereas the +proxy fence orders memory effects across proxies. Neither is a substitute for +the other. + +## Compiling and numerically verifying regular head-128 + +The regular head-128 specialization targets compute capability 10, and its +TMA/tcgen05 forms require an SM100-class GPU. The environment uses B200, CUDA +12.9 or newer, and the dependencies specified here. + +First install a CUDA-enabled PyTorch build that supports B200 using the +[official PyTorch selector](https://pytorch.org/get-started/locally/). The +`tirx-kernels` repository imports PyTorch but does not declare it as a package +dependency. + +Install TVM and `tirx-kernels` as follows: + +```bash +python -m pip install "apache-tvm==0.26.0" cuda-bindings +git clone https://github.com/mlc-ai/tirx-kernels.git +cd tirx-kernels +git checkout 5be39749e7dfd2c4bdae9b4d396f8ec35af07126 +pip install -e . +``` + +A one-row smoke test still exercises the full 128-head, top-k-2048 kernel while +keeping reference time modest: + +```python +from tirx_kernels.flashmla.sparse_prefill_head128_phase1 import run_test + +run_test( + label="tutorial_smoke", + s_q=1, + s_kv=8192, + topk=2048, + d_qk=576, + h_q=128, + h_kv=1, + d_v=512, + have_attn_sink=True, + have_topk_length=False, + seed=0, +) +print("compile, launch, and randomized reference check passed") +``` + +`run_test` covers three levels of verification: compilation checks code +generation, launch checks execution on the GPU, and the FP32 PyTorch oracle +checks output, maximum logits, and LSE with explicit tolerances. The numerical +comparison can expose head-partition and validity-bit errors beyond code +generation itself. + +The `tirx-kernels` CLI can run the registered regular-head128 configuration: + +```bash +python -m tirx_kernels.test \ + --kernel sparse_flashmla_prefill_head128_phase1 \ + --config bench_regular_dqk576_hq128_s4096_kv8192_topk2048 +``` + +Useful negative tests are just as important. Set `inject_invalid_indices=True` +to cover negative and too-large row IDs, and `have_topk_length=True` to exercise +the position predicate. Test an all-invalid row and confirm the documented +zero/-infinity/+infinity convention. Calling the TIRx dispatch entry for head-64 +and small-top-k shapes extends coverage to every prefill specialization in this +dispatch tree. Parity with the complete FlashMLA interface remains a separate +verification target. + +## Operator and specialization invariants + +1. **The cache invariant.** One `h_kv=1` latent KV row can serve multiple query + heads because key up-projection is absorbed into each query and value + up-projection is moved after core attention. The RoPE channel stays explicit, + and the QK scale remains the model's semantic scale. + +2. **The sparse-contract invariant.** Selection happens before this operator. + `indices` supplies the rows; duplicates remain duplicates; causal legality is + the caller's responsibility; and every `topk_length` lies in `[0, topk]`. + Sink and all-invalid conventions must remain identical in optimized and + reference paths. + +3. **The ownership invariant.** One two-CTA cluster owns one query row. The pair + partitions Q/output heads, selected K rows, and V features on different axes, + while CTA 0 warp 12 is the sole CTA-group MMA issuer. + +4. **The residency invariant.** $L$ means raw FP32 logits and $W$ means BF16 + unnormalized exponential weights. Large K, V, $L$, and $W$ workspaces are + single in-place tiles, while `NUM_BUFS=2` drives a barrier/phase ring. Among + data buffers, only the small packed-validity mask has two physical slots. + +5. **The handoff invariant.** Ready/done barriers transfer ownership of each + reusable segment, and proxy fences order generic and asynchronous SMEM + accesses. The sole issuer orders QK($k$) before PV($k-1$), while the threshold-6 + optimization preserves the reported-maximum semantics and the online LSE + recurrence. + +The first two invariants define the sparse-prefill operator semantics. The last +three define the regular head-128 specialization's schedule contract. Another +specialization selected by the dispatch bridge may change tile sizes, register +budgets, ownership, or barrier topology, but it must preserve the operator +semantics while defining its own schedule contract explicitly. + +## Exercises and further validation + +The regular head-128 specialization is one point in a dispatch space. The +TIRx dispatch tree also contains a head-64 phase-1 specialization and a head-128 +`d_qk=512` small-top-k specialization. Their different schedules demonstrate +how tile economics drive dispatch across that space. + +1. **Reproduce weight absorption.** Add a causal mask to the runnable absorption + proof. Confirm that MHA and MQA modes still agree, then intentionally change + the absorbed path's scale to $1/\sqrt{D_{latent}}$ and measure the error. + +2. **Stress sparse validity.** Extend `sparse_prefill_reference` with an + all-invalid query, duplicated indices, `topk_length=0`, and sink values of + both infinities. Write down the expected `(out, max_logits, lse)` for each. + +3. **Trace ownership.** For one top-k tile, label every dimension of Q, K, $L$, $W$, + V, and O with `(CTA, local row, local column)`. Show where a logical row needs + information owned by the other CTA. + +4. **Audit residency.** Starting at the `SMEMPool` allocation, draw every alias + interval. Verify why the $d_{sq}$ Q prefix remains live while the 384-column + suffix can move to TMEM, and identify the barrier that ends each reuse hazard. + +5. **Measure the threshold.** Instrument how often `should_scale_o` is true for + random and adversarial logits. Compare threshold 6 with always-rebase and + never-rebase versions in both numerical error and TMEM O traffic. + +6. **Compare dispatches.** Use `select_kernel` on head counts 64 and 128, + `d_qk` 512 and 576, and top-k values around 1280. Predict the selected module + before running it, then inspect which constraints belong to the front door and + which are enforced by an individual specialization. + +7. **Read the generated program.** Compile the smoke shape, inspect the emitted + PTX for `tcgen05` MMA, TMA gather, mbarrier, and proxy-fence instructions, and + map each one back to a TIRx line. Then run the numerical check again: source, + generated code, and observed values are three complementary kinds of proof. + +The central lesson applies beyond FlashMLA. A high-performance irregular +operator often regularizes work in stages: an indexer creates sparse addresses, +TMA gathers those addresses into dense tiles, tensor cores consume the tiles, +and explicit barriers protect aggressive storage reuse. Understanding the +algorithm, dispatch contract, ownership map, and memory protocol together is +what turns a fast kernel from an opaque artifact into an explainable program. diff --git a/img/flashmla_cache_story.png b/img/flashmla_cache_story.png new file mode 100644 index 00000000..b7abf8af Binary files /dev/null and b/img/flashmla_cache_story.png differ diff --git a/img/flashmla_cache_story_zh.svg b/img/flashmla_cache_story_zh.svg new file mode 100644 index 00000000..df752212 --- /dev/null +++ b/img/flashmla_cache_story_zh.svg @@ -0,0 +1,6165 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/flashmla_cta_ownership.png b/img/flashmla_cta_ownership.png new file mode 100644 index 00000000..24213f32 Binary files /dev/null and b/img/flashmla_cta_ownership.png differ diff --git a/img/flashmla_cta_ownership_zh.svg b/img/flashmla_cta_ownership_zh.svg new file mode 100644 index 00000000..91fc363f --- /dev/null +++ b/img/flashmla_cta_ownership_zh.svg @@ -0,0 +1,4737 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/flashmla_dataflow.png b/img/flashmla_dataflow.png new file mode 100644 index 00000000..b1b21ba8 Binary files /dev/null and b/img/flashmla_dataflow.png differ diff --git a/img/flashmla_dataflow_zh.svg b/img/flashmla_dataflow_zh.svg new file mode 100644 index 00000000..368c5bb8 --- /dev/null +++ b/img/flashmla_dataflow_zh.svg @@ -0,0 +1,5170 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/flashmla_mla_modes.png b/img/flashmla_mla_modes.png new file mode 100644 index 00000000..9d888ef1 Binary files /dev/null and b/img/flashmla_mla_modes.png differ diff --git a/img/flashmla_mla_modes_zh.svg b/img/flashmla_mla_modes_zh.svg new file mode 100644 index 00000000..129fb391 --- /dev/null +++ b/img/flashmla_mla_modes_zh.svg @@ -0,0 +1,5780 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/flashmla_pipeline.png b/img/flashmla_pipeline.png new file mode 100644 index 00000000..3686fc2e Binary files /dev/null and b/img/flashmla_pipeline.png differ diff --git a/img/flashmla_pipeline_stages.png b/img/flashmla_pipeline_stages.png new file mode 100644 index 00000000..232dc615 Binary files /dev/null and b/img/flashmla_pipeline_stages.png differ diff --git a/img/flashmla_pipeline_stages_zh.svg b/img/flashmla_pipeline_stages_zh.svg new file mode 100644 index 00000000..1dc8e413 --- /dev/null +++ b/img/flashmla_pipeline_stages_zh.svg @@ -0,0 +1,6055 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/flashmla_pipeline_zh.svg b/img/flashmla_pipeline_zh.svg new file mode 100644 index 00000000..825584b2 --- /dev/null +++ b/img/flashmla_pipeline_zh.svg @@ -0,0 +1,6205 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/flashmla_sparse_story.png b/img/flashmla_sparse_story.png new file mode 100644 index 00000000..fb3fa979 Binary files /dev/null and b/img/flashmla_sparse_story.png differ diff --git a/img/flashmla_sparse_story_zh.svg b/img/flashmla_sparse_story_zh.svg new file mode 100644 index 00000000..39e34ab3 --- /dev/null +++ b/img/flashmla_sparse_story_zh.svg @@ -0,0 +1,6671 @@ + + + + + + + + image/svg+xml + + + Modern GPU Programming for MLSys + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/scripts/flashmla_diagram_common.py b/img/scripts/flashmla_diagram_common.py new file mode 100644 index 00000000..7602a82b --- /dev/null +++ b/img/scripts/flashmla_diagram_common.py @@ -0,0 +1,228 @@ +"""Shared deterministic drawing helpers for the FlashMLA tutorial figures.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import matplotlib + + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +from matplotlib import font_manager +from matplotlib.patches import FancyArrowPatch, FancyBboxPatch, Rectangle + + +COLORS = { + "ink": "#1f2937", + "muted": "#4b5563", + "grid": "#d1d5db", + "line": "#64748b", + "cta0": "#bfdbfe", + "cta0_dark": "#2563eb", + "cta1": "#fbcfe8", + "cta1_dark": "#db2777", + "cross01": "#ddd6fe", + "cross10": "#ccfbf1", + "gmem": "#dbeafe", + "tma": "#bfdbfe", + "smem": "#e9d5ff", + "tmem": "#fed7aa", + "mma": "#bbf7d0", + "softmax": "#ddd6fe", + "barrier": "#fde68a", + "projection": "#fecaca", + "neutral": "#f8fafc", + "note": "#fefce8", +} + + +def configure_style(lang: str, font_path: str | None = None) -> None: + """Configure a fixed backend, font stack, and deterministic SVG IDs.""" + + plt.rcParams.update( + { + "font.family": "DejaVu Sans", + "font.sans-serif": ["DejaVu Sans"], + "font.size": 10, + "mathtext.fontset": "dejavusans", + "axes.unicode_minus": False, + "svg.fonttype": "none", + "svg.hashsalt": "flashmla-tutorial-v1", + } + ) + if lang != "zh": + return + + candidates = [ + font_path, + os.environ.get("FLASHMLA_CJK_FONT"), + "/tmp/NotoSansCJKsc-Regular.otf", + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", + "/usr/share/fonts/opentype/noto/NotoSansCJKsc-Regular.otf", + "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc", + ] + selected = next( + (Path(item) for item in candidates if item and Path(item).is_file()), None + ) + if selected is None: + raise FileNotFoundError( + "Chinese output requires a CJK font. Pass --font-path or set FLASHMLA_CJK_FONT." + ) + font_manager.fontManager.addfont(selected) + plt.rcParams["font.family"] = font_manager.FontProperties(fname=selected).get_name() + # Path Chinese glyphs so the checked-in SVG renders without a host CJK font. + plt.rcParams["svg.fonttype"] = "path" + + +def tr(lang: str, en: str, zh: str) -> str: + return zh if lang == "zh" else en + + +def rounded_box( + ax, + x: float, + y: float, + w: float, + h: float, + text: str, + color: str, + *, + fontsize: float = 9, + weight: str = "bold", + edgecolor: str | None = None, + linewidth: float = 1.2, + linestyle: str = "-", + zorder: int = 3, + text_color: str | None = None, +) -> FancyBboxPatch: + patch = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.04,rounding_size=0.06", + linewidth=linewidth, + linestyle=linestyle, + edgecolor=edgecolor or COLORS["ink"], + facecolor=color, + zorder=zorder, + ) + ax.add_patch(patch) + ax.text( + x + w / 2, + y + h / 2, + text, + ha="center", + va="center", + fontsize=fontsize, + weight=weight, + color=text_color or COLORS["ink"], + zorder=zorder + 1, + ) + return patch + + +def arrow( + ax, + start: tuple[float, float], + end: tuple[float, float], + *, + label: str | None = None, + color: str | None = None, + linewidth: float = 1.25, + linestyle: str = "-", + rad: float = 0.0, + label_offset: tuple[float, float] = (0.0, 0.12), + zorder: int = 2, +) -> FancyArrowPatch: + line_color = color or COLORS["line"] + patch = FancyArrowPatch( + start, + end, + arrowstyle="-|>", + mutation_scale=11, + linewidth=linewidth, + linestyle=linestyle, + color=line_color, + connectionstyle=f"arc3,rad={rad}", + zorder=zorder, + ) + ax.add_patch(patch) + if label: + mx = (start[0] + end[0]) / 2 + label_offset[0] + my = (start[1] + end[1]) / 2 + label_offset[1] + ax.text( + mx, + my, + label, + ha="center", + va="center", + fontsize=7.4, + color=line_color, + bbox=dict( + boxstyle="round,pad=0.12", + facecolor="white", + edgecolor="none", + alpha=0.92, + ), + zorder=zorder + 1, + ) + return patch + + +def plain_rect( + ax, + x: float, + y: float, + w: float, + h: float, + color: str, + *, + edgecolor: str | None = None, + linewidth: float = 1.1, + hatch: str | None = None, + zorder: int = 2, +) -> Rectangle: + patch = Rectangle( + (x, y), + w, + h, + facecolor=color, + edgecolor=edgecolor or COLORS["ink"], + linewidth=linewidth, + hatch=hatch, + zorder=zorder, + ) + ax.add_patch(patch) + return patch + + +def save_figure(fig, output: str | Path, *, dpi: int = 160) -> Path: + path = Path(output) + if path.suffix.lower() not in {".png", ".svg"}: + raise ValueError("--output must end in .png or .svg") + path.parent.mkdir(parents=True, exist_ok=True) + if path.suffix.lower() == ".svg": + metadata = {"Date": None, "Creator": "Modern GPU Programming for MLSys"} + else: + metadata = {"Software": "Modern GPU Programming for MLSys"} + fig.savefig( + path, + dpi=dpi, + facecolor="white", + transparent=False, + bbox_inches=None, + metadata=metadata, + ) + plt.close(fig) + if path.suffix.lower() == ".svg": + # Matplotlib emits trailing spaces in multiline SVG path data. Normalize + # generated assets so newly staged files pass Git's whitespace checks. + lines = path.read_text(encoding="utf-8").splitlines() + path.write_text( + "\n".join(line.rstrip() for line in lines) + "\n", + encoding="utf-8", + ) + return path diff --git a/img/scripts/gen_flashmla_cache_story.py b/img/scripts/gen_flashmla_cache_story.py new file mode 100644 index 00000000..bbbe0450 --- /dev/null +++ b/img/scripts/gen_flashmla_cache_story.py @@ -0,0 +1,459 @@ +#!/usr/bin/env python3 +"""Generate the introductory MHA-to-MLA cache story for the FlashMLA tutorial.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch + +from flashmla_diagram_common import ( + COLORS, + arrow, + configure_style, + plain_rect, + rounded_box, + save_figure, + tr, +) + + +def _panel(ax, x, y, w, h, title, subtitle, color) -> None: + panel = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.06,rounding_size=0.08", + linewidth=1.35, + edgecolor=color, + facecolor="#ffffff", + linestyle="--", + zorder=0, + ) + ax.add_patch(panel) + ax.text( + x + 0.25, + y + h - 0.34, + title, + ha="left", + va="center", + fontsize=11.5, + weight="bold", + color=color, + ) + ax.text( + x + 0.25, + y + h - 0.72, + subtitle, + ha="left", + va="center", + fontsize=7.8, + color=COLORS["muted"], + ) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(18.0, 9.0)) + ax.set_xlim(0, 18) + ax.set_ylim(0, 9) + ax.axis("off") + + ax.text( + 9, + 8.68, + tr( + lang, + "How Can One MLA Cache Entry Serve 128 Query Heads?", + "一条 MLA Cache Entry 如何服务 128 个 Query Heads？", + ), + ha="center", + va="center", + fontsize=18, + weight="bold", + ) + ax.text( + 9, + 8.28, + tr( + lang, + "store one shared compressed state; keep head-specific work on the query and output sides", + "只保存一份共享压缩状态；各 head 的差异留在 query 与 output 两侧", + ), + ha="center", + va="center", + fontsize=9.3, + color=COLORS["muted"], + ) + + _panel( + ax, + 0.35, + 1.30, + 4.45, + 6.45, + tr(lang, "1 · Ordinary MHA cache", "1 · 普通 MHA cache"), + tr( + lang, + "store a key and a value for every head", + "为每个 head 保存一份 key 和 value", + ), + COLORS["cta0_dark"], + ) + _panel( + ax, + 5.15, + 1.30, + 5.05, + 6.45, + tr(lang, "2 · MLA shared cache", "2 · MLA 共享 cache"), + tr( + lang, + "compress once and store one shared state", + "压缩一次，只保存一份共享状态", + ), + "#7c3aed", + ) + _panel( + ax, + 10.55, + 1.30, + 7.10, + 6.45, + tr(lang, "3 · Use the shared cache", "3 · 使用共享 cache"), + tr( + lang, + "head-specific work happens before and after attention", + "各 head 的特定计算放在 attention 前后", + ), + COLORS["cta1_dark"], + ) + + # MHA: a visual stack of materialized per-head cache entries. + ax.text( + 2.58, + 6.53, + tr(lang, "one cached token", "一个 cached token"), + ha="center", + fontsize=9.5, + weight="bold", + ) + row_entries = ( + ( + "head 0", + tr(lang, "cached\nkey", "缓存\nkey"), + tr(lang, "cached\nvalue", "缓存\nvalue"), + ), + ( + "head 1", + tr(lang, "cached\nkey", "缓存\nkey"), + tr(lang, "cached\nvalue", "缓存\nvalue"), + ), + ("...", "...", "..."), + ( + "head 127", + tr(lang, "cached\nkey", "缓存\nkey"), + tr(lang, "cached\nvalue", "缓存\nvalue"), + ), + ) + for idx, (label, k_label, v_label) in enumerate(row_entries): + y = 5.70 - idx * 0.83 + ax.text(0.72, y + 0.29, label, ha="left", va="center", fontsize=7.7) + plain_rect(ax, 1.55, y, 1.18, 0.58, COLORS["gmem"]) + plain_rect(ax, 2.73, y, 1.18, 0.58, "#bfdbfe") + ax.text( + 2.14, + y + 0.29, + k_label, + ha="center", + va="center", + fontsize=7.4, + weight="bold", + ) + ax.text( + 3.32, + y + 0.29, + v_label, + ha="center", + va="center", + fontsize=7.4, + weight="bold", + ) + rounded_box( + ax, + 0.82, + 1.72, + 3.50, + 0.88, + tr( + lang, + "cache stores a separate key/value slice\nfor every head", + "cache 为每个 head 保存\n独立的 key/value slice", + ), + COLORS["note"], + fontsize=8.5, + weight="normal", + edgecolor="#ca8a04", + ) + + # MLA: compress the token's content once and store it with one shared + # positional channel. Matrix names are intentionally deferred to the + # derivation that follows this introductory figure. + rounded_box( + ax, + 5.45, + 5.79, + 1.45, + 0.80, + tr(lang, "token\nrepresentation", "token\n表示"), + COLORS["neutral"], + fontsize=8.1, + ) + rounded_box( + ax, + 7.12, + 5.79, + 1.42, + 0.80, + tr(lang, "compress\ncontent once", "只压缩一次\ncontent"), + COLORS["projection"], + fontsize=8.0, + ) + rounded_box( + ax, + 8.78, + 5.69, + 1.05, + 1.00, + tr(lang, "shared\ncontent state\n512 numbers", "共享\ncontent 状态\n512 个数"), + COLORS["smem"], + fontsize=7.4, + ) + arrow(ax, (6.90, 6.19), (7.12, 6.19)) + arrow(ax, (8.54, 6.19), (8.78, 6.19)) + rounded_box( + ax, + 5.55, + 4.47, + 2.05, + 0.78, + tr(lang, "position information", "位置信息"), + "#bfdbfe", + fontsize=8.4, + ) + rounded_box( + ax, + 8.02, + 4.47, + 1.81, + 0.78, + tr(lang, "shared position part\n64 numbers", "共享位置信息\n64 个数"), + "#bfdbfe", + fontsize=7.8, + ) + arrow(ax, (7.60, 4.86), (8.02, 4.86)) + + ax.text( + 7.68, + 3.94, + tr(lang, "one cache entry per token", "每个 token 只有一条 cache entry"), + ha="center", + va="center", + fontsize=9.2, + weight="bold", + color="#6d28d9", + ) + plain_rect(ax, 5.82, 2.90, 2.76, 0.76, COLORS["gmem"]) + plain_rect(ax, 8.58, 2.90, 1.20, 0.76, "#bfdbfe") + ax.text( + 7.20, + 3.28, + tr(lang, "shared compressed\ncontent · 512", "共享压缩 content\n· 512"), + ha="center", + va="center", + fontsize=7.8, + weight="bold", + ) + ax.text( + 9.18, + 3.28, + tr(lang, "position\npart · 64", "位置信息\n· 64"), + ha="center", + va="center", + fontsize=7.5, + weight="bold", + ) + rounded_box( + ax, + 5.82, + 1.72, + 3.96, + 0.72, + tr( + lang, + "one shared entry · not 128 key/value pairs", + "一份共享 cache entry，而不是 128 组 key/value", + ), + COLORS["note"], + fontsize=8.2, + weight="normal", + edgecolor="#ca8a04", + ) + arrow(ax, (9.78, 3.28), (11.15, 4.09), color="#7c3aed", linewidth=1.6, rad=-0.05) + + # Head-specific query/output work surrounds attention over the shared cache. + rounded_box( + ax, + 10.92, + 5.89, + 1.46, + 0.70, + tr( + lang, + "content part of query\nfor one head", + "某个 head query 的\ncontent 部分", + ), + COLORS["neutral"], + fontsize=7.8, + ) + rounded_box( + ax, + 12.78, + 5.89, + 1.66, + 0.70, + tr(lang, "head-specific\nquery transform", "该 head 的\nquery 变换"), + COLORS["projection"], + fontsize=7.7, + ) + rounded_box( + ax, + 10.92, + 4.93, + 1.46, + 0.70, + tr(lang, "position part of query\nfor that head", "该 head query 的\n位置信息"), + "#bfdbfe", + fontsize=7.7, + ) + rounded_box( + ax, + 14.88, + 5.34, + 2.28, + 1.02, + tr( + lang, + "query used by attention\ncontent view + position", + "attention 使用的 query\ncontent 视图 + 位置信息", + ), + COLORS["tmem"], + fontsize=7.8, + ) + arrow(ax, (12.38, 6.24), (12.78, 6.24)) + arrow(ax, (14.44, 6.24), (14.88, 6.03)) + arrow( + ax, + (12.38, 5.28), + (14.88, 5.62), + label=tr(lang, "position stays separate", "位置信息保持独立"), + color=COLORS["cta0_dark"], + rad=-0.05, + label_offset=(0.0, -0.14), + ) + + rounded_box( + ax, + 11.15, + 3.63, + 2.56, + 0.92, + tr( + lang, + "one shared cache entry\nkey side: content + position\nvalue side: content", + "一条共享 cache entry\nkey 侧：压缩 content + 位置信息\nvalue 侧：压缩 content", + ), + COLORS["gmem"], + fontsize=7.6, + ) + rounded_box( + ax, + 14.15, + 3.48, + 2.12, + 1.10, + tr(lang, "attention over\nthe shared cache", "在共享 cache 上\n计算 attention"), + COLORS["mma"], + fontsize=8.7, + ) + arrow(ax, (16.02, 5.34), (15.56, 4.58), color=COLORS["line"], rad=0.04) + arrow(ax, (13.71, 4.09), (14.15, 4.03)) + rounded_box( + ax, + 13.22, + 2.22, + 1.60, + 0.72, + tr(lang, "shared-space result\n512 numbers", "共享空间结果\n512 个数"), + COLORS["tmem"], + fontsize=7.6, + ) + rounded_box( + ax, + 15.22, + 2.22, + 1.90, + 0.72, + tr( + lang, + "head-specific\noutput transform", + "该 head 的\noutput 变换", + ), + COLORS["projection"], + fontsize=7.6, + ) + arrow(ax, (15.05, 3.48), (14.10, 2.94), color=COLORS["line"], rad=0.04) + arrow(ax, (14.82, 2.58), (15.22, 2.58)) + ax.text( + 15.18, + 1.74, + tr( + lang, + "repeat the query/output work for all 128 heads\n→ 128 distinct head outputs", + "对全部 128 个 heads 重复 query/output 处理\n→ 128 份不同的 head output", + ), + ha="center", + va="center", + fontsize=8.0, + color=COLORS["muted"], + weight="bold", + ) + + rounded_box( + ax, + 1.65, + 0.30, + 14.70, + 0.66, + tr( + lang, + "MLA caches one shared compressed state per token. Head-specific transformations happen around attention, so the cache does not store 128 key/value pairs.", + "MLA 为每个 token 只缓存一份共享压缩状态。各 head 的特定变换放在 attention 两侧，因此 cache 无需保存 128 组 key/value。", + ), + COLORS["note"], + fontsize=8.4, + weight="normal", + edgecolor="#ca8a04", + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/img/scripts/gen_flashmla_cta_ownership.py b/img/scripts/gen_flashmla_cta_ownership.py new file mode 100644 index 00000000..db3e9b81 --- /dev/null +++ b/img/scripts/gen_flashmla_cta_ownership.py @@ -0,0 +1,339 @@ +#!/usr/bin/env python3 +"""Generate the sparse FlashMLA head128 regular 2-CTA ownership diagram.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt + +from flashmla_diagram_common import ( + COLORS, + arrow, + configure_style, + plain_rect, + rounded_box, + save_figure, + tr, +) + + +def _partition_bar(ax, x, y, w, h, *, labels, colors, split="vertical"): + if split == "vertical": + cell_w = w / len(labels) + for idx, (label, color) in enumerate(zip(labels, colors, strict=True)): + plain_rect(ax, x + idx * cell_w, y, cell_w, h, color) + ax.text( + x + (idx + 0.5) * cell_w, + y + h / 2, + label, + ha="center", + va="center", + fontsize=8.6, + weight="bold", + color=COLORS["ink"], + zorder=4, + ) + else: + cell_h = h / len(labels) + for idx, (label, color) in enumerate(zip(labels, colors, strict=True)): + plain_rect(ax, x, y + (len(labels) - idx - 1) * cell_h, w, cell_h, color) + ax.text( + x + w / 2, + y + (len(labels) - idx - 0.5) * cell_h, + label, + ha="center", + va="center", + fontsize=8.6, + weight="bold", + color=COLORS["ink"], + zorder=4, + ) + + +def _logical_matrix(ax, x, y, w, h, *, title, col_axis, lang, kind): + ax.text( + x + w / 2, + y + h + 0.38, + title, + ha="center", + va="center", + fontsize=11, + weight="bold", + ) + ax.text( + x + w / 2, + y - 0.3, + col_axis, + ha="center", + va="center", + fontsize=8, + color=COLORS["muted"], + ) + ax.text( + x - 0.34, + y + h / 2, + ( + tr(lang, "Q-head rows", "Q head 行") + if kind == "k" + else tr(lang, "W rows (Q-head ownership)", "W 行（沿用 Q head 所有权）") + ), + ha="center", + va="center", + rotation=90, + fontsize=8, + color=COLORS["muted"], + ) + cell_colors = [ + [COLORS["cta0"], COLORS["cross01"]], + [COLORS["cross10"], COLORS["cta1"]], + ] + q_labels = ["Q0", "Q1"] if kind == "k" else ["W(Q0)", "W(Q1)"] + b_labels = ["K0", "K1"] if kind == "k" else ["V0", "V1"] + for row in range(2): + for col in range(2): + cx = x + col * w / 2 + cy = y + (1 - row) * h / 2 + plain_rect(ax, cx, cy, w / 2, h / 2, cell_colors[row][col]) + ax.text( + cx + w / 4, + cy + h / 4, + f"{q_labels[row]} × {b_labels[col]}", + ha="center", + va="center", + fontsize=8.5 if kind == "v" else 9, + weight="bold", + color=COLORS["ink"], + zorder=4, + ) + ax.text( + x + 0.10, + y + h - 0.11, + "heads 0:64", + ha="left", + va="top", + fontsize=6.7, + color=COLORS["muted"], + zorder=5, + ) + ax.text( + x + 0.10, + y + h / 2 - 0.11, + "heads 64:128", + ha="left", + va="top", + fontsize=6.7, + color=COLORS["muted"], + zorder=5, + ) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(16.0, 8.8)) + ax.set_xlim(0, 16) + ax.set_ylim(0, 8.8) + ax.axis("off") + + ax.text( + 8, + 8.5, + tr( + lang, + "Sparse FlashMLA Head128 Regular: 2-CTA Ownership", + "Sparse FlashMLA Head128 Regular：2-CTA 所有权", + ), + ha="center", + va="center", + fontsize=18, + weight="bold", + ) + ax.text( + 8, + 8.12, + tr( + lang, + "one CTA pair per query row · B_H=128 · B_TOPK=128 · D_V=512", + "每个 query row 使用一个 CTA pair · B_H=128 · B_TOPK=128 · D_V=512", + ), + ha="center", + va="center", + fontsize=9.5, + color=COLORS["muted"], + ) + + headings = [ + (0.45, tr(lang, "Q ownership", "Q 所有权")), + (5.5, tr(lang, "K selected-token ownership", "K 的选中 token 所有权")), + (10.55, tr(lang, "V feature ownership", "V feature 所有权")), + ] + for x, text in headings: + ax.text( + x, + 7.66, + text, + ha="left", + va="center", + fontsize=11, + weight="bold", + color=COLORS["ink"], + ) + + _partition_bar( + ax, + 0.65, + 6.13, + 4.15, + 1.15, + labels=("CTA 0 · Q heads 0:64", "CTA 1 · Q heads 64:128"), + colors=(COLORS["cta0"], COLORS["cta1"]), + split="horizontal", + ) + ax.text( + 2.72, + 5.88, + tr(lang, "each CTA: 64 × d_qk", "每个 CTA：64 × d_qk"), + ha="center", + fontsize=8, + color=COLORS["muted"], + ) + + _partition_bar( + ax, + 5.7, + 6.13, + 4.15, + 1.15, + labels=("CTA 0\nslots 0:64", "CTA 1\nslots 64:128"), + colors=(COLORS["cta0"], COLORS["cta1"]), + ) + ax.text( + 7.77, + 5.88, + tr(lang, "within the current sparse B_TOPK tile", "当前 sparse B_TOPK tile 内"), + ha="center", + fontsize=8, + color=COLORS["muted"], + ) + + _partition_bar( + ax, + 10.75, + 6.13, + 4.15, + 1.15, + labels=("CTA 0\nfeatures 0:256", "CTA 1\nfeatures 256:512"), + colors=(COLORS["cta0"], COLORS["cta1"]), + ) + ax.text( + 12.82, + 5.88, + tr( + lang, + "each half covers all 128 selected-token rows", + "每个 feature half 覆盖全部 128 个选中 token 行", + ), + ha="center", + fontsize=8, + color=COLORS["muted"], + ) + + _logical_matrix( + ax, + 0.85, + 1.75, + 5.3, + 2.7, + title=tr( + lang, + "QK · cta_group::2 → logits L [128 heads × 128 slots]", + "QK · cta_group::2 → logits L [128 heads × 128 slots]", + ), + col_axis=tr(lang, "selected-token columns", "选中 token 列"), + lang=lang, + kind="k", + ) + _logical_matrix( + ax, + 9.85, + 1.75, + 5.3, + 2.7, + title=tr( + lang, + "PV · cta_group::2 → output O [128 heads × 512 features]", + "PV · cta_group::2 → output O [128 heads × 512 features]", + ), + col_axis=tr(lang, "output-feature columns", "输出 feature 列"), + lang=lang, + kind="v", + ) + + rounded_box( + ax, + 6.70, + 2.48, + 2.6, + 1.25, + tr( + lang, + "L → WG0 softmax → W\nsame head × slot shape", + "L → WG0 softmax → W\n保持相同的 head × slot 形状", + ), + COLORS["softmax"], + fontsize=8.7, + ) + arrow(ax, (6.15, 3.1), (6.70, 3.1), label="L") + arrow(ax, (9.30, 3.1), (9.85, 3.1), label="W") + # Route ownership guides around the logical-matrix titles, then terminate + # on the tile boundaries. A direct downward arrow would cross the long + # QK/PV headings and make the target ambiguous at the inline HTML width. + guide_specs = ( + ((2.72, 5.76), 0.42, (0.85, 4.34), COLORS["cta0_dark"]), + ((7.77, 5.76), 6.55, (6.15, 4.34), COLORS["cta1_dark"]), + ((12.82, 5.76), 15.58, (15.15, 4.34), COLORS["cta1_dark"]), + ) + for start, corridor_x, target, color in guide_specs: + ax.plot( + [start[0], corridor_x, corridor_x], + [start[1], start[1], 4.53], + color=color, + linewidth=1.0, + zorder=1, + ) + arrow( + ax, + (corridor_x, 4.53), + target, + color=color, + linewidth=1.0, + zorder=2, + ) + + rounded_box( + ax, + 1.0, + 0.52, + 14.0, + 0.62, + tr( + lang, + "Every colored quadrant is a required block product. Cross-colored quadrants are not cross-CTA reductions; cta_group::2 combines operand halves without an explicit DSMEM copy.", + "每个彩色象限都是必需的 block product。交叉配色象限不是跨 CTA reduction；cta_group::2 在没有显式 DSMEM copy 的情况下组合 operand halves。", + ), + COLORS["note"], + fontsize=8.7, + weight="normal", + edgecolor="#ca8a04", + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/img/scripts/gen_flashmla_dataflow.py b/img/scripts/gen_flashmla_dataflow.py new file mode 100644 index 00000000..1534affa --- /dev/null +++ b/img/scripts/gen_flashmla_dataflow.py @@ -0,0 +1,510 @@ +#!/usr/bin/env python3 +"""Generate the sparse FlashMLA head128 regular SMEM/TMEM data-flow diagram.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch + +from flashmla_diagram_common import ( + COLORS, + arrow as _arrow, + configure_style, + plain_rect, + rounded_box as _rounded_box, + save_figure, + tr, +) + + +def box(ax, x, y, w, h, text, color, *, fontsize=11.0, **kwargs): + """Draw a node sized for the book's roughly 790 px content column.""" + + return _rounded_box( + ax, + x, + y, + w, + h, + text, + color, + fontsize=fontsize, + **kwargs, + ) + + +def edge( + ax, + start, + end, + *, + label=None, + label_offset=(0.0, 0.14), + fontsize=10.0, + **kwargs, +): + """Draw an edge and place its label in surrounding whitespace.""" + + line_color = kwargs.get("color") or COLORS["line"] + zorder = kwargs.get("zorder", 2) + patch = _arrow(ax, start, end, label=None, **kwargs) + if label: + ax.text( + (start[0] + end[0]) / 2 + label_offset[0], + (start[1] + end[1]) / 2 + label_offset[1], + label, + ha="center", + va="center", + fontsize=fontsize, + color=line_color, + bbox=dict( + boxstyle="round,pad=0.10", + facecolor="white", + edgecolor="none", + alpha=0.94, + ), + zorder=zorder + 1, + ) + return patch + + +def panel(ax, x, y, w, h, title, subtitle): + """Draw one readable layer of the single data-flow figure.""" + + patch = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.05,rounding_size=0.06", + linewidth=1.15, + linestyle="--", + edgecolor="#c4b5fd", + facecolor="#ffffff", + zorder=0, + ) + ax.add_patch(patch) + ax.text( + x + 0.20, + y + h - 0.24, + title, + ha="left", + va="center", + fontsize=13.0, + weight="bold", + color="#6d28d9", + ) + ax.text( + x + w - 0.20, + y + h - 0.24, + subtitle, + ha="right", + va="center", + fontsize=10.0, + color=COLORS["muted"], + ) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(12.0, 13.0)) + ax.set_xlim(0, 14) + ax.set_ylim(0, 13) + ax.axis("off") + + ax.text( + 7, + 12.62, + tr( + lang, + "Sparse FlashMLA Head128 Regular: Data Residency", + "Sparse FlashMLA Head128 Regular：数据驻留", + ), + ha="center", + va="center", + fontsize=19, + weight="bold", + ) + ax.text( + 7, + 12.25, + tr( + lang, + "one CTA's logical view; cta_group::2 combines the pair's operand halves", + "单个 CTA 的 logical view；cta_group::2 组合 CTA pair 的 operand halves", + ), + ha="center", + va="center", + fontsize=11.0, + color=COLORS["muted"], + ) + + panel( + ax, + 0.25, + 8.20, + 13.50, + 3.55, + tr(lang, "1 · QK and online softmax", "1 · QK 与 online softmax"), + tr(lang, "BF16 operands → FP32 L → BF16 W", "BF16 operands → FP32 L → BF16 W"), + ) + + box(ax, 0.48, 9.20, 0.78, 0.92, "Q\nGMEM", COLORS["gmem"], fontsize=12.0) + box( + ax, + 1.52, + 9.08, + 1.55, + 1.16, + tr( + lang, + "q_full\nSMEM · TMA\n64×d_qk / CTA", + "q_full\nSMEM · TMA\n64×d_qk / CTA", + ), + COLORS["smem"], + fontsize=10.5, + ) + box( + ax, + 3.35, + 9.86, + 1.90, + 0.80, + "Q prefix\nSMEM\nd_sq columns", + COLORS["smem"], + fontsize=10.2, + ) + box( + ax, + 3.35, + 8.66, + 1.90, + 0.80, + tr(lang, "Q suffix · TMEM\n384 columns", "Q suffix · TMEM\n384 列"), + COLORS["tmem"], + fontsize=10.2, + ) + box( + ax, + 5.57, + 9.01, + 1.42, + 1.30, + "QK MMA\nSS prefix\n+ TS suffix", + COLORS["mma"], + fontsize=10.5, + ) + box( + ax, + 7.37, + 9.01, + 1.42, + 1.30, + "L (tmem_p)\nFP32 · TMEM\n64×128", + COLORS["tmem"], + fontsize=9.5, + ) + box( + ax, + 9.17, + 8.93, + 1.90, + 1.46, + tr( + lang, + "WG0 softmax\nmask · max\nexp · row sum", + "WG0 softmax\nmask · max\nexp · row sum", + ), + COLORS["softmax"], + fontsize=10.5, + ) + box( + ax, + 11.42, + 9.83, + 2.05, + 0.78, + tr( + lang, + "max_logits / lse\nGMEM\nsink excluded", + "max_logits / lse\nGMEM\n不含 sink", + ), + COLORS["gmem"], + fontsize=9.5, + ) + box( + ax, + 11.42, + 8.69, + 2.05, + 0.78, + "mi / real_mi / li\nWG0 registers", + COLORS["neutral"], + fontsize=10.0, + ) + + edge(ax, (1.26, 9.66), (1.52, 9.66)) + edge(ax, (3.07, 9.82), (3.35, 10.26)) + edge(ax, (3.07, 9.48), (3.35, 9.06)) + edge(ax, (5.25, 10.26), (5.57, 9.98)) + edge(ax, (5.25, 9.06), (5.57, 9.34)) + edge(ax, (6.99, 9.66), (7.37, 9.66)) + edge(ax, (8.79, 9.66), (9.17, 9.66)) + edge(ax, (11.07, 9.86), (11.42, 10.22)) + edge(ax, (11.07, 9.40), (11.42, 9.08)) + + panel( + ax, + 0.25, + 3.20, + 13.50, + 4.70, + tr( + lang, + "2 · Sparse gather, PV, and epilogue", + "2 · Sparse gather、PV 与 epilogue", + ), + tr( + lang, + "indices + KV feed both gather4 paths; O~ / (li + sink)", + "indices + KV 同时送入两条 gather4；O~ / (li + sink)", + ), + ) + + box(ax, 0.52, 6.00, 1.05, 0.88, "indices\nGMEM", COLORS["gmem"], fontsize=11.5) + box( + ax, + 0.52, + 4.00, + 1.15, + 1.05, + "KV cache\nGMEM\n[s_kv,1,\nd_qk]", + COLORS["gmem"], + fontsize=8.8, + ) + box( + ax, + 2.15, + 5.90, + 2.05, + 1.08, + tr( + lang, + "WG1 K · SMEM\ngather4\n64×d_qk / CTA", + "WG1 K · SMEM\ngather4\n64×d_qk / CTA", + ), + COLORS["smem"], + fontsize=10.2, + ) + box( + ax, + 2.15, + 4.00, + 2.05, + 1.05, + tr( + lang, + "WG2 V · SMEM\ngather4\n128×256 / CTA", + "WG2 V · SMEM\ngather4\n128×256 / CTA", + ), + COLORS["smem"], + fontsize=10.2, + ) + box( + ax, + 4.55, + 5.90, + 1.85, + 1.08, + tr( + lang, + "validity · SMEM\nbounds + length\n2 × 16 bytes", + "validity · SMEM\nbounds + length\n2 × 16 bytes", + ), + COLORS["barrier"], + fontsize=10.0, + ) + box( + ax, + 6.55, + 5.90, + 1.95, + 1.08, + "W · BF16\ns_smem_gemm\n64 × 128", + COLORS["smem"], + fontsize=9.4, + ) + box(ax, 6.82, 4.00, 1.40, 1.05, "PV MMA\nW × V", COLORS["mma"], fontsize=11.4) + box( + ax, + 8.62, + 4.00, + 1.42, + 1.05, + "O~ · FP32\nTMEM\n64×512", + COLORS["tmem"], + fontsize=9.8, + ) + box( + ax, + 10.30, + 3.90, + 1.75, + 1.25, + tr( + lang, + "WG0 epilogue\nO~/(li + sink)\n→ BF16", + "WG0 epilogue\nO~/(li + sink)\n→ BF16", + ), + COLORS["softmax"], + fontsize=9.5, + ) + box( + ax, + 8.85, + 5.95, + 1.55, + 0.95, + tr( + lang, + "attn_sink\nGMEM", + "attn_sink\nGMEM", + ), + COLORS["gmem"], + fontsize=9.5, + ) + box( + ax, + 12.34, + 3.96, + 1.20, + 1.13, + tr( + lang, + "o_smem\nBF16\n64×512\nTMA → out", + "o_smem\nBF16\n64×512\nTMA → out", + ), + COLORS["smem"], + fontsize=9.2, + ) + + # Both gather producers consume both the sparse row coordinates and KV source. + # The junction makes that all-to-all relationship explicit without four labels. + shared = (1.82, 5.48) + edge(ax, (1.57, 6.25), shared) + edge(ax, (1.67, 4.78), shared) + edge(ax, shared, (2.15, 6.25)) + edge(ax, shared, (2.15, 4.78)) + ax.plot(*shared, marker="o", markersize=4.0, color=COLORS["line"], zorder=4) + edge(ax, (4.20, 6.44), (4.55, 6.44)) + + # K enters the QK panel; validity and W cross the layer boundary separately. + ax.plot([4.20, 4.38, 6.28], [6.56, 8.02, 8.02], color=COLORS["line"], lw=1.25) + edge(ax, (6.28, 8.02), (6.28, 9.01)) + ax.plot([6.40, 6.40, 10.10], [6.62, 7.98, 7.98], color=COLORS["line"], lw=1.25) + edge(ax, (10.10, 7.98), (10.10, 8.93)) + ax.plot([9.60, 9.60, 7.52], [8.93, 8.08, 8.08], color=COLORS["line"], lw=1.25) + edge(ax, (7.52, 8.08), (7.52, 6.98)) + edge(ax, (4.20, 4.52), (6.82, 4.52)) + edge(ax, (7.52, 5.90), (7.52, 5.05)) + edge(ax, (8.22, 4.52), (8.62, 4.52)) + edge(ax, (10.04, 4.52), (10.30, 4.52)) + edge(ax, (10.40, 6.28), (10.88, 5.15)) + edge(ax, (12.05, 4.52), (12.34, 4.52)) + + panel( + ax, + 0.25, + 0.35, + 13.50, + 2.45, + tr( + lang, + "3 · Per-CTA SMEM lifetime aliasing", + "3 · 每个 CTA 的 SMEM 生命周期复用", + ), + tr( + lang, + "one in-place K/V region; o_smem later reuses its base", + "单份原位 K/V 区域；随后 o_smem 复用其基址", + ), + ) + box( + ax, + 0.55, + 0.77, + 2.10, + 0.92, + "q_full · SMEM\nprefix | suffix384", + COLORS["smem"], + fontsize=10.2, + ) + edge(ax, (2.65, 1.23), (3.08, 1.23)) + + plain_rect(ax, 3.08, 0.77, 1.45, 0.92, COLORS["smem"], edgecolor=COLORS["ink"]) + plain_rect(ax, 4.53, 0.77, 1.75, 0.92, "#ddd6fe", edgecolor=COLORS["ink"]) + plain_rect(ax, 6.28, 0.77, 1.75, 0.92, "#c4b5fd", edgecolor=COLORS["ink"]) + ax.text( + 3.80, + 1.23, + tr(lang, "Q prefix\nlive", "Q prefix\n存活"), + ha="center", + va="center", + fontsize=11.0, + weight="bold", + ) + ax.text( + 5.40, + 1.23, + tr(lang, "one V\nworkspace", "单份 V\nworkspace"), + ha="center", + va="center", + fontsize=11.0, + weight="bold", + ) + ax.text( + 7.15, + 1.23, + tr(lang, "one K\nworkspace", "单份 K\nworkspace"), + ha="center", + va="center", + fontsize=11.0, + weight="bold", + ) + + edge(ax, (8.03, 1.23), (8.46, 1.23)) + box( + ax, + 8.46, + 0.77, + 1.95, + 0.92, + "o_smem · BF16\n64×512", + COLORS["smem"], + fontsize=10.2, + ) + ax.text( + 10.82, + 1.23, + tr( + lang, + "Adjacent K/V workspaces\nbegin at Q's released suffix.", + "相邻的 K/V workspace\n从已释放的 Q suffix 开始。", + ), + ha="left", + va="center", + fontsize=11.0, + color=COLORS["muted"], + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/img/scripts/gen_flashmla_mla_modes.py b/img/scripts/gen_flashmla_mla_modes.py new file mode 100644 index 00000000..538a8b41 --- /dev/null +++ b/img/scripts/gen_flashmla_mla_modes.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +"""Generate the two equivalent MLA attention-core views used by the tutorial.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch + +from flashmla_diagram_common import ( + COLORS, + arrow, + configure_style, + plain_rect, + rounded_box, + save_figure, + tr, +) + + +def _panel( + ax, x: float, y: float, w: float, h: float, title: str, subtitle: str, color: str +) -> None: + panel = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.06,rounding_size=0.08", + linewidth=1.35, + edgecolor=color, + facecolor="#ffffff", + linestyle="--", + zorder=0, + ) + ax.add_patch(panel) + ax.text( + x + 0.28, + y + h - 0.34, + title, + ha="left", + va="center", + fontsize=12, + weight="bold", + color=color, + ) + ax.text( + x + 0.28, + y + h - 0.72, + subtitle, + ha="left", + va="center", + fontsize=7.8, + color=COLORS["muted"], + ) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(18.0, 9.4)) + ax.set_xlim(0, 18) + ax.set_ylim(0, 9.4) + ax.axis("off") + + ax.text( + 9, + 9.08, + tr( + lang, + "MLA: Two Equivalent Attention-Core Views of One Cache", + "MLA：同一份 Cache 的两种等价 Attention Core 视图", + ), + ha="center", + va="center", + fontsize=18, + weight="bold", + ) + ax.text( + 9, + 8.70, + tr( + lang, + "move the up-projections across the attention core without changing the cached latent state", + "在 attention core 两侧移动 up-projection，而不改变缓存的 latent state", + ), + ha="center", + va="center", + fontsize=9.3, + color=COLORS["muted"], + ) + + # One common compressed cache is the source for both views. + ax.text( + 9, + 8.26, + tr(lang, "one MLA KV cache entry", "同一条 MLA KV cache entry"), + ha="center", + va="center", + fontsize=10, + weight="bold", + ) + plain_rect(ax, 6.15, 7.35, 4.25, 0.72, COLORS["gmem"], edgecolor=COLORS["ink"]) + plain_rect(ax, 10.40, 7.35, 1.45, 0.72, "#bfdbfe", edgecolor=COLORS["ink"]) + ax.text( + 8.27, + 7.71, + "c_KV [latent 512]", + ha="center", + va="center", + fontsize=9.3, + weight="bold", + zorder=4, + ) + ax.text( + 11.12, + 7.71, + "k_R [RoPE 64]", + ha="center", + va="center", + fontsize=8.8, + weight="bold", + zorder=4, + ) + ax.text( + 9, + 7.15, + tr( + lang, + "stored once; the diagrams below differ only in where W_UK and W_UV run", + "只缓存一份；下方两种视图只改变 W_UK 与 W_UV 的执行位置", + ), + ha="center", + va="center", + fontsize=7.8, + color=COLORS["muted"], + ) + + _panel( + ax, + 0.35, + 1.42, + 8.32, + 5.30, + tr(lang, "MHA mode", "MHA mode"), + tr( + lang, + "expand per-head K and V before the core", + "在 core 之前展开每个 head 的 K 与 V", + ), + COLORS["cta0_dark"], + ) + _panel( + ax, + 9.33, + 1.42, + 8.32, + 5.30, + tr(lang, "Absorbed MQA mode", "Absorbed MQA mode"), + tr( + lang, + "absorb W_UK into q_C; q_R bypasses; move W_UV after the core", + "将 W_UK 吸收到 q_C；q_R 绕过；把 W_UV 移到 core 之后", + ), + COLORS["cta1_dark"], + ) + arrow(ax, (7.52, 7.35), (4.58, 6.43), color=COLORS["cta0_dark"], rad=0.08) + arrow(ax, (10.47, 7.35), (13.55, 6.43), color=COLORS["cta1_dark"], rad=-0.08) + + # Left: explicit per-head K/V expansion before a conventional MHA-shaped core. + rounded_box( + ax, + 0.82, + 4.82, + 2.05, + 0.82, + "per-head Q\n[q_C 128 ; q_R 64]", + COLORS["neutral"], + fontsize=8.3, + ) + rounded_box( + ax, 0.82, 3.78, 1.72, 0.75, "c_KV\nlatent 512", COLORS["gmem"], fontsize=8.4 + ) + rounded_box( + ax, 0.82, 2.76, 1.72, 0.70, "k_R\nRoPE 64", COLORS["gmem"], fontsize=8.3 + ) + rounded_box(ax, 3.05, 4.02, 1.28, 0.70, "W_UK", COLORS["projection"], fontsize=9) + rounded_box(ax, 3.05, 3.00, 1.28, 0.70, "W_UV", COLORS["projection"], fontsize=9) + rounded_box( + ax, + 4.86, + 3.97, + 2.20, + 0.82, + "per-head K\n[k_C 128 ; k_R 64]", + COLORS["smem"], + fontsize=8.2, + ) + rounded_box( + ax, 4.86, 2.92, 2.20, 0.76, "per-head V\nv_C 128", COLORS["smem"], fontsize=8.4 + ) + rounded_box( + ax, + 3.08, + 1.72, + 2.70, + 0.84, + tr( + lang, + "MHA attention core\nper-head K / V", + "MHA attention core\nper-head K / V", + ), + COLORS["mma"], + fontsize=8.7, + ) + rounded_box( + ax, 6.34, 1.76, 1.75, 0.76, "per-head out\n128", COLORS["neutral"], fontsize=8.5 + ) + arrow(ax, (2.54, 4.15), (3.05, 4.37)) + arrow(ax, (2.54, 4.00), (3.05, 3.35)) + arrow(ax, (4.33, 4.37), (4.86, 4.37)) + # k_R bypasses W_UK and joins K below that projection box. + ax.plot( + [2.54, 2.76, 4.55], [3.11, 3.88, 3.88], color=COLORS["line"], lw=1.25, zorder=2 + ) + arrow( + ax, + (4.55, 3.88), + (5.15, 3.97), + label=tr(lang, "concat", "拼接"), + color=COLORS["line"], + rad=-0.03, + ) + arrow(ax, (4.33, 3.35), (4.86, 3.30)) + # Route Q around the left edge so it does not cut through c_KV/W_UV. + ax.plot( + [1.84, 0.62, 0.62, 2.72], + [4.82, 4.82, 2.14, 2.14], + color=COLORS["line"], + lw=1.25, + zorder=2, + ) + arrow(ax, (2.72, 2.14), (3.08, 2.14), color=COLORS["line"]) + arrow(ax, (5.70, 3.97), (4.82, 2.56), color=COLORS["line"], rad=0.04) + arrow(ax, (5.70, 2.92), (5.18, 2.56), color=COLORS["line"], rad=-0.02) + arrow(ax, (5.78, 2.14), (6.34, 2.14)) + + # Right: only the content query absorbs W_UK. The per-head RoPE query + # bypasses that projection and is concatenated immediately before QK. + rounded_box( + ax, + 9.78, + 5.18, + 1.48, + 0.66, + "q_C · per head\ncontent 128", + COLORS["neutral"], + fontsize=8.0, + ) + rounded_box( + ax, + 9.78, + 4.37, + 1.48, + 0.62, + "q_R · per head\nRoPE 64", + "#bfdbfe", + fontsize=8.0, + ) + rounded_box( + ax, + 11.72, + 5.14, + 1.45, + 0.74, + tr(lang, "(W_UK)^T q_C\ncontent only", "(W_UK)^T q_C\n仅 content"), + COLORS["projection"], + fontsize=7.9, + ) + rounded_box( + ax, + 13.62, + 4.64, + 2.78, + 1.02, + "absorbed Q\n[q_abs 512 ; q_R 64] = 576", + COLORS["tmem"], + fontsize=8.2, + ) + rounded_box( + ax, + 9.83, + 3.30, + 2.30, + 0.94, + "shared KV · h_kv=1\n[c_KV 512 ; k_R 64]\nV uses c_KV 512", + COLORS["gmem"], + fontsize=8.0, + ) + rounded_box( + ax, + 12.35, + 2.44, + 2.42, + 0.94, + tr(lang, "absorbed MQA\nattention core", "absorbed MQA\nattention core"), + COLORS["mma"], + fontsize=8.8, + ) + rounded_box( + ax, 15.18, 2.51, 1.55, 0.80, "latent out\n512", COLORS["tmem"], fontsize=8.5 + ) + rounded_box( + ax, + 14.30, + 1.58, + 1.45, + 0.70, + tr(lang, "W_UV\nafter core", "W_UV\n移到 core 后"), + COLORS["projection"], + fontsize=8.0, + ) + rounded_box( + ax, 16.18, 1.58, 1.05, 0.70, "out\n128", COLORS["neutral"], fontsize=8.3 + ) + arrow(ax, (11.26, 5.51), (11.72, 5.51)) + arrow(ax, (13.17, 5.51), (13.62, 5.30)) + # The explicit positional channel never passes through W_UK. + arrow( + ax, + (11.26, 4.68), + (13.62, 4.91), + label=tr(lang, "bypass W_UK", "绕过 W_UK"), + color=COLORS["cta0_dark"], + rad=-0.06, + label_offset=(0.0, -0.15), + ) + arrow(ax, (15.05, 4.64), (13.83, 3.38), color=COLORS["line"], rad=-0.05) + arrow(ax, (10.98, 3.30), (12.64, 3.08), color=COLORS["line"], rad=0.03) + arrow(ax, (14.77, 2.91), (15.18, 2.91)) + arrow(ax, (15.95, 2.51), (15.02, 2.28), color=COLORS["line"], rad=0.06) + arrow(ax, (15.75, 1.93), (16.18, 1.93)) + + # Explicit equivalence and chapter scope. + ax.text( + 9, + 1.23, + tr( + lang, + "Equivalent after projection reassociation: both produce the same per-head output 128", + "重新结合 projection 后等价：两种视图都产生相同的 per-head output 128", + ), + ha="center", + va="center", + fontsize=8.0, + color="#7c3aed", + weight="bold", + ) + rounded_box( + ax, + 2.1, + 0.24, + 13.8, + 0.76, + tr( + lang, + "The d_qk=576 sparse-prefill case studied here uses the absorbed MQA view:\nshared KV [512+64], absorbed Q 576, latent output 512. These dimensions do not describe every FlashMLA kernel.", + "本章重点分析的 d_qk=576 sparse-prefill case 采用 absorbed MQA 视图：\nshared KV [512+64]、absorbed Q 576、latent output 512。这些维度并不代表所有 FlashMLA kernel。", + ), + COLORS["note"], + fontsize=8.2, + weight="normal", + edgecolor="#ca8a04", + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/img/scripts/gen_flashmla_pipeline.py b/img/scripts/gen_flashmla_pipeline.py new file mode 100644 index 00000000..38ddeefe --- /dev/null +++ b/img/scripts/gen_flashmla_pipeline.py @@ -0,0 +1,451 @@ +#!/usr/bin/env python3 +"""Generate the sparse FlashMLA head128 regular steady-state pipeline diagram.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch + +from flashmla_diagram_common import ( + COLORS, + arrow as _arrow, + configure_style, + rounded_box as _rounded_box, + save_figure, + tr, +) + + +def box(ax, x, y, w, h, text, color, *, fontsize=10.8, **kwargs): + """Draw a timeline node sized for the book's content column.""" + + return _rounded_box( + ax, + x, + y, + w, + h, + text, + color, + fontsize=fontsize, + **kwargs, + ) + + +def edge( + ax, + start, + end, + *, + label=None, + label_offset=(0.0, 0.14), + fontsize=10.0, + **kwargs, +): + """Draw a dependency edge with an HTML-readable label.""" + + line_color = kwargs.get("color") or COLORS["line"] + zorder = kwargs.get("zorder", 2) + patch = _arrow(ax, start, end, label=None, **kwargs) + if label: + ax.text( + (start[0] + end[0]) / 2 + label_offset[0], + (start[1] + end[1]) / 2 + label_offset[1], + label, + ha="center", + va="center", + fontsize=fontsize, + color=line_color, + bbox=dict( + boxstyle="round,pad=0.10", + facecolor="white", + edgecolor="none", + alpha=0.94, + ), + zorder=zorder + 1, + ) + return patch + + +def panel(ax, x, y, w, h, title, subtitle): + """Draw one layer of the steady-state dependency story.""" + + patch = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.05,rounding_size=0.06", + linewidth=1.15, + linestyle="--", + edgecolor="#c4b5fd", + facecolor="#ffffff", + zorder=0, + ) + ax.add_patch(patch) + ax.text( + x + 0.20, + y + h - 0.24, + title, + ha="left", + va="center", + fontsize=13.0, + weight="bold", + color="#6d28d9", + ) + ax.text( + x + w - 0.20, + y + h - 0.24, + subtitle, + ha="right", + va="center", + fontsize=10.0, + color=COLORS["muted"], + ) + + +def lane(ax, y, title, role): + """Draw a named warpgroup lane and its baseline.""" + + box( + ax, + 0.48, + y - 0.38, + 1.42, + 0.76, + f"{title}\n{role}", + COLORS["neutral"], + fontsize=10.2, + ) + ax.plot([2.08, 13.45], [y, y], color="#e5e7eb", lw=1.0, zorder=0) + + +def tag(ax, x, y, text, *, fontsize=9.8): + """Place a completion/ready barrier name away from node text.""" + + ax.text( + x, + y, + text, + ha="center", + va="center", + fontsize=fontsize, + color="#92400e", + bbox=dict( + boxstyle="round,pad=0.10", + facecolor=COLORS["barrier"], + edgecolor="#d97706", + ), + zorder=7, + ) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(14.0, 13.2)) + ax.set_xlim(0, 14) + ax.set_ylim(0, 13.2) + ax.axis("off") + + ax.text( + 7, + 12.84, + tr( + lang, + "Sparse FlashMLA Head128 Regular: Steady-State Pipeline", + "Sparse FlashMLA Head128 Regular：稳态 Pipeline", + ), + ha="center", + va="center", + fontsize=20, + weight="bold", + ) + ax.text( + 7, + 12.30, + tr( + lang, + "softmax(k−1) may overlap QK(k); once QK(k) completes,\nsoftmax(k) may overlap asynchronous PV(k−1) · box widths are not cycle measurements", + "softmax(k−1) 可与 QK(k) 重叠；QK(k) 完成后，\nsoftmax(k) 可与异步 PV(k−1) 重叠 · 方框宽度不表示 cycle 数", + ), + ha="center", + va="center", + fontsize=11.5, + color=COLORS["muted"], + linespacing=1.25, + ) + + panel( + ax, + 0.25, + 6.55, + 13.50, + 5.25, + tr( + lang, + "A · One issuer order; four independent part handshakes", + "A · 唯一 issuer 的顺序与四个独立分段握手", + ), + tr(lang, "box widths are not cycle measurements", "方框宽度不表示 cycle 数"), + ) + + ax.text( + 1.05, + 10.72, + tr( + lang, + "WG3 · CTA0 warp 12\nsole MMA issuer", + "WG3 · CTA0 warp 12\n唯一 MMA issuer", + ), + ha="center", + va="center", + fontsize=10.5, + weight="bold", + color=COLORS["ink"], + ) + issuer_boxes = [ + (2.25, "wait p_free[k−1]\nQK(k) · SS prefix"), + (5.10, "QK(k) · TS suffix384"), + (7.95, "wait so_ready[k−1]\nPV(k−1) · part0"), + (10.80, "PV(k−1) · part1"), + ] + for x, text in issuer_boxes: + box(ax, x, 10.25, 2.25, 0.90, text, COLORS["mma"], fontsize=10.4) + for x0, x1 in ((4.50, 5.10), (7.35, 7.95), (10.20, 10.80)): + edge(ax, (x0, 10.70), (x1, 10.70), color=COLORS["ink"]) + ax.text( + 7.0, + 9.93, + tr(lang, "strict issue order", "严格串行发起顺序"), + ha="center", + va="center", + fontsize=10.0, + color=COLORS["muted"], + style="italic", + ) + + handshakes = [ + ( + 0.70, + 8.33, + tr( + lang, + "WG1: k_part0_ready[k]\n→ QK SS prefix\n→ qk_part_done[k]\n→ gather K(k+1) part0", + "WG1：k_part0_ready[k]\n→ QK SS prefix\n→ qk_part_done[k]\n→ gather K(k+1) part0", + ), + "#dbeafe", + ), + ( + 7.10, + 8.33, + tr( + lang, + "WG1: k_part1_ready[k]\n→ QK TS suffix384\n→ qk_done[k]\n→ gather K(k+1) part1", + "WG1：k_part1_ready[k]\n→ QK TS suffix384\n→ qk_done[k]\n→ gather K(k+1) part1", + ), + "#dbeafe", + ), + ( + 0.70, + 6.88, + tr( + lang, + "WG2: v_part0_ready[k−1]\n→ PV part0\n→ sv_part_done[k−1]\n→ gather V(k) part0", + "WG2：v_part0_ready[k−1]\n→ PV part0\n→ sv_part_done[k−1]\n→ gather V(k) part0", + ), + COLORS["tma"], + ), + ( + 7.10, + 6.88, + tr( + lang, + "WG2: v_part1_ready[k−1]\n→ PV part1\n→ sv_done[k−1]\n→ gather V(k) part1", + "WG2：v_part1_ready[k−1]\n→ PV part1\n→ sv_done[k−1]\n→ gather V(k) part1", + ), + COLORS["tma"], + ), + ] + for x, y, text, color in handshakes: + box(ax, x, y, 6.20, 1.12, text, color, fontsize=10.3, weight="normal") + + panel( + ax, + 0.25, + 1.05, + 13.50, + 5.15, + tr( + lang, + "B · Tile k: mask-slot reuse and WG0 handoff", + "B · Tile k：mask slot 复用与 WG0 交接", + ), + tr( + lang, + "arrows name the barrier that guards reuse", + "箭头标出保护复用的 barrier", + ), + ) + box( + ax, + 0.70, + 4.35, + 2.45, + 0.92, + tr(lang, "WG3 warp 13\npack mask(k)", "WG3 warp 13\npack mask(k)"), + COLORS["barrier"], + fontsize=10.8, + ) + box( + ax, + 5.35, + 4.35, + 3.10, + 0.92, + tr( + lang, + "WG0 consumes mask(k)\nwhile processing L(k)", + "WG0 处理 L(k) 时\n读取 mask(k)", + ), + COLORS["softmax"], + fontsize=10.8, + ) + box( + ax, + 10.35, + 4.35, + 2.95, + 0.92, + "WG3 warp 13\npack mask(k+2)\nreuse slot k%2", + COLORS["barrier"], + fontsize=10.6, + ) + edge( + ax, + (3.15, 4.81), + (5.35, 4.81), + label="k_valid_ready[k]", + label_offset=(0.0, 0.22), + fontsize=9.7, + color="#d97706", + ) + edge( + ax, + (8.45, 4.81), + (10.35, 4.81), + label="k_valid_free[k]", + label_offset=(0.0, 0.22), + fontsize=9.7, + color="#d97706", + linestyle="--", + ) + + box( + ax, + 0.70, + 2.35, + 2.45, + 1.00, + tr(lang, "qk_done[k]\nload L(k) from TMEM", "qk_done[k]\n从 TMEM 读取 L(k)"), + COLORS["tmem"], + fontsize=10.8, + ) + box( + ax, + 5.05, + 2.35, + 3.25, + 1.00, + tr(lang, "WG0: mask · max\nexp · row sum", "WG0：mask · max\nexp · row sum"), + COLORS["softmax"], + fontsize=10.8, + ) + box( + ax, + 9.65, + 2.20, + 3.65, + 1.30, + tr( + lang, + "wait sv_done(k−1)\nwrite W(k) · optional O rescale\narrive so_ready[k]", + "等待 sv_done(k−1)\n写 W(k) · 按需重缩放 O\narrive so_ready[k]", + ), + COLORS["softmax"], + fontsize=10.5, + ) + edge( + ax, + (3.15, 2.85), + (5.05, 2.85), + color=COLORS["line"], + ) + edge( + ax, + (8.30, 2.85), + (9.65, 2.85), + color="#7c3aed", + ) + ax.text( + 2.0, + 1.65, + "p_free[k] → QK(k+1)", + ha="center", + va="center", + fontsize=10.3, + color="#7c3aed", + weight="bold", + ) + edge( + ax, + (1.95, 2.35), + (1.95, 1.87), + color="#7c3aed", + linestyle="--", + ) + ax.text( + 11.5, + 1.65, + "so_ready[k] → PV(k)", + ha="center", + va="center", + fontsize=10.3, + color="#7c3aed", + weight="bold", + ) + edge( + ax, + (11.5, 2.20), + (11.5, 1.87), + color="#7c3aed", + linestyle="--", + ) + + box( + ax, + 0.52, + 0.10, + 12.96, + 0.72, + tr( + lang, + "NUM_BUFS=2 is a barrier/phase ring: slot=k%2, phase=(k//2)&1.\nK, V, and W use one in-place workspace each; only the small validity mask has two data slots.", + "NUM_BUFS=2 是 barrier/phase ring：slot=k%2，phase=(k//2)&1。\nK、V、W 各自只有一份原位 workspace；只有小型 validity mask 有两个 data slots。", + ), + COLORS["note"], + fontsize=9.8, + weight="normal", + edgecolor="#ca8a04", + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/img/scripts/gen_flashmla_pipeline_stages.py b/img/scripts/gen_flashmla_pipeline_stages.py new file mode 100644 index 00000000..8b5ea583 --- /dev/null +++ b/img/scripts/gen_flashmla_pipeline_stages.py @@ -0,0 +1,478 @@ +#!/usr/bin/env python3 +"""Generate the fill, steady-state, and drain story for sparse FlashMLA.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch + +from flashmla_diagram_common import ( + COLORS, + arrow, + configure_style, + rounded_box, + save_figure, + tr, +) + + +def _panel(ax, x, y, w, h, title, subtitle, color) -> None: + panel = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.06,rounding_size=0.08", + linewidth=1.35, + edgecolor=color, + facecolor="#ffffff", + linestyle="--", + zorder=0, + ) + ax.add_patch(panel) + ax.text( + x + 0.24, + y + h - 0.32, + title, + ha="left", + va="center", + fontsize=11.2, + weight="bold", + color=color, + ) + ax.text( + x + 0.24, + y + h - 0.70, + subtitle, + ha="left", + va="center", + fontsize=7.6, + color=COLORS["muted"], + ) + + +def _lane(ax, x0, x1, y, label, lang) -> None: + ax.text( + x0, + y, + label, + ha="left", + va="center", + fontsize=7.1, + color=COLORS["muted"], + weight="bold", + ) + ax.plot([x0 + 0.76, x1], [y, y], color="#e5e7eb", lw=1.0, zorder=0) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(18.0, 9.6)) + ax.set_xlim(0, 18) + ax.set_ylim(0, 9.6) + ax.axis("off") + + ax.text( + 9, + 9.27, + tr( + lang, + "How Does the One-Step-Lag Pipeline Fill and Drain?", + "相差一个 Tile 的 Pipeline 如何填充与排空？", + ), + ha="center", + va="center", + fontsize=18, + weight="bold", + ) + ax.text( + 9, + 8.88, + tr( + lang, + "N = num_k_blocks; one WG3 issuer loop executes k=0…N and serially issues QK before PV", + "N = num_k_blocks；一个 WG3 issuer loop 执行 k=0…N，并按顺序先发 QK、再发 PV", + ), + ha="center", + va="center", + fontsize=9.2, + color=COLORS["muted"], + ) + + # The mathematical dependency is simple; scheduling staggers adjacent tiles. + rounded_box( + ax, + 2.02, + 7.83, + 2.02, + 0.62, + "QK(t) → L(t)", + COLORS["mma"], + fontsize=8.6, + ) + rounded_box( + ax, + 5.35, + 7.83, + 2.60, + 0.62, + "softmax(t) → W(t)", + COLORS["softmax"], + fontsize=8.6, + ) + rounded_box( + ax, + 9.24, + 7.83, + 2.30, + 0.62, + "PV(t) → O~(t)", + COLORS["mma"], + fontsize=8.6, + ) + arrow(ax, (4.04, 8.14), (5.35, 8.14), label="L(t)") + arrow(ax, (7.95, 8.14), (9.24, 8.14), label="W(t)") + ax.text( + 13.95, + 8.14, + tr( + lang, + "logical dependency for one tile\nexecution overlaps neighboring tiles", + "单个 tile 的逻辑依赖\n执行时与相邻 tiles 重叠", + ), + ha="center", + va="center", + fontsize=8.1, + color=COLORS["muted"], + style="italic", + ) + + y0 = 1.42 + h = 5.94 + _panel( + ax, + 0.35, + y0, + 4.40, + h, + tr(lang, "Fill · k=0", "填充 · k=0"), + tr( + lang, + "the issuer has no previous tile to multiply", + "issuer 尚无上一 tile 可执行 PV", + ), + COLORS["cta0_dark"], + ) + _panel( + ax, + 5.02, + y0, + 8.33, + h, + tr(lang, "Steady state · 1 ≤ k < N", "稳态 · 1 ≤ k < N"), + tr( + lang, + "adjacent softmax work overlaps asynchronous MMAs", + "相邻 tile 的 softmax 与异步 MMA 重叠", + ), + "#7c3aed", + ) + _panel( + ax, + 13.62, + y0, + 4.03, + h, + tr(lang, "Drain · k=N", "排空 · k=N"), + tr(lang, "the issuer has no new QK tile", "issuer 不再有新的 QK tile"), + COLORS["cta1_dark"], + ) + + # Fill: Q prologue, QK(0), and the first softmax; no PV is issued at k=0. + rounded_box( + ax, + 0.72, + 5.82, + 3.66, + 0.58, + tr( + lang, + "prologue: Q TMA · TMEM alloc · initial K/V gather", + "prologue：Q TMA · TMEM alloc · 首轮 K/V gather", + ), + COLORS["gmem"], + fontsize=7.6, + weight="normal", + ) + _lane(ax, 0.68, 4.38, 5.25, "WG3", lang) + _lane(ax, 0.68, 4.38, 4.08, "WG0", lang) + _lane(ax, 0.68, 4.38, 2.87, "WG1/2", lang) + rounded_box(ax, 1.65, 4.88, 1.56, 0.72, "QK(0)", COLORS["mma"], fontsize=8.8) + rounded_box( + ax, + 3.47, + 4.91, + 0.66, + 0.66, + "PV\n—", + COLORS["neutral"], + fontsize=7.7, + weight="normal", + edgecolor="#9ca3af", + linestyle="--", + ) + rounded_box( + ax, + 2.10, + 3.71, + 1.82, + 0.74, + "softmax(0)", + COLORS["softmax"], + fontsize=8.4, + ) + arrow(ax, (2.84, 4.88), (2.72, 4.45), label="L(0)", rad=0.03) + rounded_box( + ax, + 1.37, + 2.50, + 2.65, + 0.72, + tr(lang, "prepare safe next segments", "准备下一批可安全复用的 segments"), + COLORS["tma"], + fontsize=7.8, + weight="normal", + ) + ax.text( + 2.55, + 1.91, + tr(lang, "issuer branch: QK(0) only", "issuer 分支：仅 QK(0)"), + ha="center", + va="center", + fontsize=8.1, + color=COLORS["muted"], + weight="bold", + ) + + # Steady state: make the two legal overlaps explicit without depicting two issuers. + _lane(ax, 5.34, 13.00, 5.65, "WG3", lang) + _lane(ax, 5.34, 13.00, 4.20, "WG0", lang) + _lane(ax, 5.34, 13.00, 2.72, "WG1/2", lang) + arrow(ax, (6.15, 6.37), (12.78, 6.37), color="#9ca3af", linewidth=1.0) + ax.text( + 9.46, + 6.44, + tr(lang, "issuer program order / time", "issuer 程序顺序 / 时间"), + ha="center", + va="bottom", + fontsize=7.2, + color=COLORS["muted"], + style="italic", + ) + + # Light overlap bands sit behind the actual operation boxes. + ax.axvspan(6.14, 8.58, ymin=0.41, ymax=0.66, color="#eff6ff", alpha=0.75, zorder=-1) + ax.axvspan( + 9.62, 12.47, ymin=0.41, ymax=0.66, color="#fdf2f8", alpha=0.75, zorder=-1 + ) + rounded_box(ax, 6.38, 5.27, 1.86, 0.76, "QK(k)", COLORS["mma"], fontsize=9.0) + rounded_box( + ax, + 9.72, + 5.27, + 2.16, + 0.76, + "PV(k−1)", + COLORS["mma"], + fontsize=9.0, + ) + arrow( + ax, + (8.24, 5.65), + (9.72, 5.65), + label=tr(lang, "same issuer · serial issue", "同一 issuer · 串行发出"), + color=COLORS["ink"], + linewidth=1.5, + label_offset=(0.0, 0.18), + ) + rounded_box( + ax, + 6.05, + 3.82, + 2.58, + 0.76, + "softmax(k−1) → W(k−1)", + COLORS["softmax"], + fontsize=8.0, + ) + rounded_box( + ax, + 9.88, + 3.82, + 2.52, + 0.76, + "softmax(k) → W(k)", + COLORS["softmax"], + fontsize=8.0, + ) + arrow( + ax, + (8.63, 4.20), + (10.03, 5.27), + label="W(k−1) ready", + color="#7c3aed", + rad=-0.06, + label_offset=(0.35, 0.25), + ) + arrow( + ax, + (7.93, 5.27), + (10.12, 4.58), + label="L(k) ready", + color="#7c3aed", + rad=0.05, + label_offset=(0.40, -0.34), + ) + ax.text( + 7.34, + 3.43, + tr(lang, "legal overlap A", "合法重叠 A"), + ha="center", + va="center", + fontsize=7.4, + color=COLORS["cta0_dark"], + weight="bold", + ) + ax.text( + 11.11, + 3.43, + tr(lang, "legal overlap B", "合法重叠 B"), + ha="center", + va="center", + fontsize=7.4, + color=COLORS["cta1_dark"], + weight="bold", + ) + rounded_box( + ax, + 6.02, + 2.35, + 2.82, + 0.72, + tr(lang, "gather safe K(k+1) parts", "gather 可复用的 K(k+1) parts"), + COLORS["tma"], + fontsize=7.7, + weight="normal", + ) + rounded_box( + ax, + 9.55, + 2.35, + 2.82, + 0.72, + tr(lang, "gather safe V(k) parts", "gather 可复用的 V(k) parts"), + COLORS["tma"], + fontsize=7.7, + weight="normal", + ) + ax.text( + 9.18, + 1.87, + tr( + lang, + "QK(k) and PV(k−1) are not concurrently issued", + "QK(k) 与 PV(k−1) 并非同时发出", + ), + ha="center", + va="center", + fontsize=8.2, + color="#6d28d9", + weight="bold", + ) + + # Drain: the extra loop iteration performs only the final PV, then WG0 normalizes. + _lane(ax, 13.92, 17.30, 5.65, "WG3", lang) + _lane(ax, 13.92, 17.30, 4.20, "WG0", lang) + _lane(ax, 13.92, 17.30, 2.72, "WG1/2", lang) + rounded_box( + ax, + 14.64, + 5.30, + 0.86, + 0.70, + "QK\n—", + COLORS["neutral"], + fontsize=7.6, + weight="normal", + edgecolor="#9ca3af", + linestyle="--", + ) + rounded_box(ax, 15.78, 5.27, 1.34, 0.76, "PV(N−1)", COLORS["mma"], fontsize=8.1) + arrow(ax, (15.50, 5.65), (15.78, 5.65), color=COLORS["ink"]) + rounded_box( + ax, + 14.53, + 3.77, + 2.58, + 0.86, + tr( + lang, + "wait sv_done(N−1)\nnormalize + epilogue", + "等待 sv_done(N−1)\n归一化 + epilogue", + ), + COLORS["softmax"], + fontsize=7.8, + ) + arrow(ax, (16.48, 5.27), (16.20, 4.63), color="#7c3aed", rad=0.03) + rounded_box( + ax, + 14.62, + 2.37, + 2.42, + 0.70, + tr(lang, "no more gathers", "不再发起 gather"), + COLORS["neutral"], + fontsize=8.0, + weight="normal", + edgecolor="#9ca3af", + linestyle="--", + ) + ax.text( + 15.64, + 1.87, + tr(lang, "issuer branch: PV(N−1) only", "issuer 分支：仅 PV(N−1)"), + ha="center", + va="center", + fontsize=8.0, + color=COLORS["muted"], + weight="bold", + ) + + rounded_box( + ax, + 1.25, + 0.30, + 15.50, + 0.72, + tr( + lang, + "Source-order loop: for k in [0, N]: if k < N issue QK(k); if k > 0 issue PV(k−1). MMAs execute asynchronously, but one warp issues them in that order.", + "源码顺序：for k in [0, N]：若 k < N 则发 QK(k)；若 k > 0 则发 PV(k−1)。MMA 异步执行，但由同一个 warp 按此顺序发出。", + ), + COLORS["note"], + fontsize=8.15, + weight="normal", + edgecolor="#ca8a04", + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/img/scripts/gen_flashmla_sparse_story.py b/img/scripts/gen_flashmla_sparse_story.py new file mode 100644 index 00000000..3ad3d617 --- /dev/null +++ b/img/scripts/gen_flashmla_sparse_story.py @@ -0,0 +1,467 @@ +#!/usr/bin/env python3 +"""Generate the selection-versus-attention story for sparse FlashMLA prefill.""" + +from __future__ import annotations + +import argparse + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch + +from flashmla_diagram_common import ( + COLORS, + arrow, + configure_style, + plain_rect, + rounded_box, + save_figure, + tr, +) + + +def _panel(ax, x, y, w, h, title, subtitle, color) -> None: + panel = FancyBboxPatch( + (x, y), + w, + h, + boxstyle="round,pad=0.06,rounding_size=0.08", + linewidth=1.35, + edgecolor=color, + facecolor="#ffffff", + linestyle="--", + zorder=0, + ) + ax.add_patch(panel) + ax.text( + x + 0.25, + y + h - 0.34, + title, + ha="left", + va="center", + fontsize=11.3, + weight="bold", + color=color, + ) + ax.text( + x + 0.25, + y + h - 0.72, + subtitle, + ha="left", + va="center", + fontsize=7.8, + color=COLORS["muted"], + ) + + +def draw(lang: str, output: str, font_path: str | None = None) -> None: + configure_style(lang, font_path) + fig, ax = plt.subplots(figsize=(18.0, 9.2)) + ax.set_xlim(0, 18) + ax.set_ylim(0, 9.2) + ax.axis("off") + + ax.text( + 9, + 8.86, + tr( + lang, + "Who Chooses the Tokens in Sparse FlashMLA?", + "Sparse FlashMLA 中由谁选择 token？", + ), + ha="center", + va="center", + fontsize=18, + weight="bold", + ) + ax.text( + 9, + 8.46, + tr( + lang, + "the indexer ranks candidates; sparse prefill receives row addresses and only performs attention", + "indexer 对 candidates 排序；sparse prefill 接收 row addresses，只负责 attention", + ), + ha="center", + va="center", + fontsize=9.3, + color=COLORS["muted"], + ) + + _panel( + ax, + 0.35, + 1.34, + 4.55, + 6.65, + tr(lang, "1 · Select outside the kernel", "1 · 在 kernel 外完成选择"), + tr( + lang, + "lightning indexer scores candidate rows", + "lightning indexer 为 candidate rows 打分", + ), + COLORS["cta0_dark"], + ) + _panel( + ax, + 5.18, + 1.34, + 5.80, + 6.65, + tr(lang, "2 · Pass an address list", "2 · 传入 address list"), + tr( + lang, + "one list is shared by all 128 query heads", + "同一份 list 由全部 128 个 query heads 共享", + ), + "#7c3aed", + ) + _panel( + ax, + 11.26, + 1.34, + 6.39, + 6.65, + tr(lang, "3 · Attend over selected rows", "3 · 对选中 rows 执行 attention"), + tr( + lang, + "gather first; then run the dense tile chain", + "先 gather，再执行规则的 tile 计算链", + ), + COLORS["cta1_dark"], + ) + + # Selection is an upstream operator, not part of this sparse-prefill kernel. + rounded_box( + ax, + 0.75, + 6.10, + 1.26, + 0.74, + "query q", + COLORS["neutral"], + fontsize=8.7, + ) + rounded_box( + ax, + 0.75, + 4.65, + 1.26, + 0.90, + tr(lang, "candidate\nKV rows", "candidate\nKV rows"), + COLORS["gmem"], + fontsize=8.4, + ) + rounded_box( + ax, + 2.40, + 5.12, + 1.87, + 1.30, + tr( + lang, + "lightning indexer\nscore + top-k", + "lightning indexer\nscore + top-k", + ), + "#fef3c7", + fontsize=9.0, + edgecolor="#d97706", + ) + arrow(ax, (2.01, 6.47), (2.40, 6.05)) + arrow(ax, (2.01, 5.10), (2.40, 5.55)) + rounded_box( + ax, + 1.03, + 3.32, + 3.18, + 0.90, + tr( + lang, + "output: selected row addresses\nnot gathered K/V values", + "输出：选中 rows 的 addresses\n不是已 gather 的 K/V values", + ), + COLORS["barrier"], + fontsize=8.4, + weight="normal", + ) + arrow(ax, (3.34, 5.12), (2.91, 4.22)) + rounded_box( + ax, + 0.78, + 1.78, + 3.68, + 0.78, + tr( + lang, + "different operator · its own cost", + "独立 operator · 有自己的计算成本", + ), + COLORS["note"], + fontsize=8.4, + weight="normal", + edgecolor="#ca8a04", + ) + + # Concrete semantics for one query row. + ax.text( + 8.08, + 6.67, + "indices[q, 0, :]", + ha="center", + va="center", + fontsize=10, + weight="bold", + ) + values = ("42", "7", "42", "−1", "91", "13") + states = ( + tr(lang, "valid", "有效"), + tr(lang, "valid", "有效"), + tr(lang, "duplicate", "重复"), + "OOB", + tr(lang, "valid", "有效"), + tr(lang, "length tail", "长度 tail"), + ) + fills = ( + COLORS["cta0"], + COLORS["cta0"], + COLORS["cross01"], + "#fecaca", + COLORS["cta0"], + "#e5e7eb", + ) + x0 = 5.60 + cell_w = 0.82 + for idx, (value, state, fill) in enumerate(zip(values, states, fills, strict=True)): + x = x0 + idx * cell_w + plain_rect(ax, x, 5.82, cell_w, 0.64, fill) + ax.text( + x + cell_w / 2, + 6.14, + value, + ha="center", + va="center", + fontsize=9, + weight="bold", + ) + ax.text( + x + cell_w / 2, + 5.56, + state, + ha="center", + va="center", + fontsize=6.5, + color=COLORS["muted"], + ) + ax.plot([x0, x0 + 5 * cell_w], [5.30, 5.30], color="#7c3aed", lw=1.3) + ax.plot([x0, x0], [5.24, 5.36], color="#7c3aed", lw=1.3) + ax.plot([x0 + 5 * cell_w, x0 + 5 * cell_w], [5.24, 5.36], color="#7c3aed", lw=1.3) + ax.text( + x0 + 2.5 * cell_w, + 5.10, + "topk_length[q] = 5", + ha="center", + va="center", + fontsize=7.8, + color="#6d28d9", + weight="bold", + ) + + rounded_box( + ax, + 5.70, + 4.05, + 4.82, + 0.76, + tr( + lang, + "valid = 0 ≤ idx < s_kv AND slot < topk_length", + "valid = 0 ≤ idx < s_kv 且 slot < topk_length", + ), + COLORS["barrier"], + fontsize=8.3, + weight="normal", + ) + rounded_box( + ax, + 5.82, + 2.93, + 2.18, + 0.70, + "h_kv = 1", + COLORS["gmem"], + fontsize=8.8, + ) + rounded_box( + ax, + 8.32, + 2.93, + 2.06, + 0.70, + "Q heads 0…127", + COLORS["neutral"], + fontsize=8.3, + ) + arrow( + ax, + (8.00, 3.28), + (8.32, 3.28), + label=tr(lang, "same list", "同一份 list"), + color="#7c3aed", + label_offset=(0.0, 0.17), + ) + ax.text( + 8.08, + 2.20, + tr( + lang, + "duplicates participate twice; invalid slots are masked", + "重复 row 会参与两次；invalid slots 被 mask", + ), + ha="center", + va="center", + fontsize=8.1, + color=COLORS["muted"], + weight="bold", + ) + + # The kernel turns irregular addresses into regular QK/softmax/PV tiles. + rounded_box( + ax, + 11.66, + 6.11, + 1.46, + 0.74, + "Q\n[128, d_qk]", + COLORS["neutral"], + fontsize=8.0, + ) + rounded_box( + ax, + 11.66, + 4.72, + 1.46, + 0.92, + "KV cache\n[s_kv, 1, d_qk]", + COLORS["gmem"], + fontsize=7.7, + ) + rounded_box( + ax, + 13.51, + 4.94, + 1.32, + 1.46, + tr( + lang, + "gather\nselected rows\n+ validity", + "gather\n选中 rows\n+ validity", + ), + COLORS["smem"], + fontsize=7.8, + ) + arrow(ax, (13.12, 5.18), (13.51, 5.40), label="rows", rad=-0.05) + arrow( + ax, + (10.98, 5.86), + (13.51, 6.07), + label="indices", + color="#7c3aed", + rad=-0.06, + ) + + rounded_box(ax, 15.25, 5.95, 1.04, 0.72, "QK", COLORS["mma"], fontsize=9.0) + rounded_box( + ax, + 16.55, + 5.95, + 0.78, + 0.72, + "L", + COLORS["tmem"], + fontsize=9.0, + ) + arrow(ax, (13.12, 6.48), (15.25, 6.31), label="Q") + arrow(ax, (14.83, 5.84), (15.25, 6.13), label="K", rad=-0.05) + arrow(ax, (16.29, 6.31), (16.55, 6.31)) + + rounded_box( + ax, + 15.02, + 4.48, + 2.10, + 0.92, + tr( + lang, + "mask L → softmax\n→ BF16 weights W", + "mask L → softmax\n→ BF16 weights W", + ), + COLORS["softmax"], + fontsize=8.0, + ) + arrow(ax, (16.94, 5.95), (16.25, 5.40), label="L", rad=0.05) + rounded_box(ax, 14.15, 3.02, 1.16, 0.76, "PV\nW × V", COLORS["mma"], fontsize=8.5) + arrow(ax, (15.65, 4.48), (14.98, 3.78), label="W", rad=0.04) + arrow( + ax, + (14.18, 4.94), + (14.47, 3.78), + label=tr(lang, "V = first 512 cols", "V = 前 512 cols"), + color=COLORS["line"], + rad=-0.12, + label_offset=(-0.60, -0.02), + ) + rounded_box( + ax, + 15.72, + 2.82, + 1.52, + 1.16, + "out [128,512]\nmax_logits\nlse", + COLORS["gmem"], + fontsize=7.8, + ) + arrow(ax, (15.31, 3.40), (15.72, 3.40)) + + rounded_box( + ax, + 11.80, + 1.76, + 5.32, + 0.66, + tr( + lang, + "the kernel performs no ranking; it consumes addresses and computes attention", + "kernel 不执行排序；它消费 addresses 并计算 attention", + ), + COLORS["note"], + fontsize=8.2, + weight="normal", + edgecolor="#ca8a04", + ) + + rounded_box( + ax, + 1.55, + 0.30, + 14.90, + 0.66, + tr( + lang, + "No causal flag exists in this interface: the caller's list determines allowed keys; the kernel only applies bounds and optional length masking.", + "该接口没有 causal flag：允许访问哪些 keys 由 caller 的 list 决定；kernel 只应用边界检查与可选的 length mask。", + ), + COLORS["note"], + fontsize=8.3, + weight="normal", + edgecolor="#ca8a04", + ) + + save_figure(fig, output, dpi=160) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lang", choices=("en", "zh"), default="en") + parser.add_argument("--output", required=True) + parser.add_argument("--font-path") + args = parser.parse_args() + draw(args.lang, args.output, args.font_path) diff --git a/index.md b/index.md index 6975f80b..45ca9fac 100644 --- a/index.md +++ b/index.md @@ -14,9 +14,9 @@ high-performance implementation. This book develops both. The book proceeds from hardware to programming model to complete kernels. It first introduces GPU organization and execution, then presents the programming model used throughout the book, and finally builds high-performance kernels step by step. The main target is NVIDIA Blackwell, and the -running examples are General Matrix-Matrix Multiplication (GEMM) and FlashAttention. Along the way, -the book develops the key ideas behind GPU optimization: data layout, asynchronous data movement, -and asynchronous coordination. +running examples are General Matrix-Matrix Multiplication (GEMM), FlashAttention, and FlashMLA. +Along the way, the book develops the key ideas behind GPU optimization: data layout, asynchronous +data movement, and asynchronous coordination. The material grows out of the [Machine Learning Systems](https://mlsyscourse.org/) course series at Carnegie Mellon University. The examples use the **TIRx** Python DSL so that the ideas can be @@ -37,8 +37,9 @@ This book is open source. Contributions, corrections, and examples are welcome t foundation for the code examples throughout the book. - **Part III, GEMM: Tiled to SOTA.** A complete guide to optimizing a tiled GEMM, built up through TMA pipelining, persistent scheduling, warp specialization, and 2-CTA clusters. -- **Part IV, Flash Attention 4.** A complete attention kernel built from the Part III techniques: - two MMAs with softmax between them, online-softmax rescaling, causal masking, and GQA. +- **Part IV, Attention Kernels.** This part first builds Flash Attention 4 from the Part III + techniques, then introduces Multi-head Latent Attention and maps a sparse FlashMLA prefill + operator to a two-CTA Blackwell kernel. - **Appendices.** TIRx language reference, a reproducible GPU benchmarking and profiling workflow, compiler internals, and a guide to debugging asynchronous kernels. @@ -75,10 +76,11 @@ chapter_gemm_advanced/index ``` ```{toctree} -:caption: Part IV, Flash Attention 4 +:caption: Part IV, Attention Kernels :maxdepth: 2 chapter_flash_attention/index +chapter_flashmla/index ``` ```{toctree} diff --git a/zh/chapter_flashmla/index.md b/zh/chapter_flashmla/index.md new file mode 100644 index 00000000..d9abf62b --- /dev/null +++ b/zh/chapter_flashmla/index.md @@ -0,0 +1,898 @@ +(chap_flashmla)= +# FlashMLA + +:::{admonition} 概览 +:class: overview + +- 从普通 MHA 的 KV cache 出发，理解 MLA 为什么只需为每个 token 保存一份共享的压缩状态，以及各 head 特有的 K/V 变换被移到了哪里。 +- 理解 FlashMLA sparse attention 的输入：外部 indexer 先选出 KV rows，kernel 再对这些 rows 完成 QK、softmax 和 PV；随后用一个可执行 reference 明确其数值语义和边界条件。 +- 以一个具体实现为例，理解 sparse attention 如何映射到 Blackwell，并在 B200 上完成编译与数值验证。 +::: + +语言模型生成文本时，通常一次只预测一个新 token。为了完成这一步，attention 会用当前位置的 query 与所有历史 tokens 的 key 和 value 计算相关性，再根据 attention weights 汇总历史信息。如果不保存历史 K/V，模型就必须在每一步重新计算整个 prefix 的 K/V，造成大量重复计算。因此，系统会把已经得到的 K/V 保存在 **KV cache** 中，供后续 attention 直接读取。生成模型通常采用 causal attention，位置 $s$ 的 K/V 只依赖位置 $s$ 及其之前的 tokens，后来追加的 token 不会改变它们。Prefill 会并行计算 prompt 中的 tokens 并填充 cache；decode 随后逐步读取已有 cache，再追加新 token 的 K/V。 + +普通 multi-head attention（MHA）中，每个 attention head 都有自己的 K 和 V，因此 KV cache 必须为每个已经处理过的 token 保存各个 head 的 K/V。上下文越长，cache 占用的显存越多；生成每个新 token 时，attention 还要读取不断增长的历史 K/V，使它同时面临显存容量和内存带宽的压力。**Multi-head Latent Attention（MLA）** 通过改变 cache 表示来减少这些开销：它把每个 token 中与 K/V 相关的信息压缩成一份由所有 heads 共享的状态。计算 attention 时，每个 head 仍使用自己的变换，因此能够保持各自的 attention 行为。 + +FlashMLA 是 DeepSeek 为 MLA 开发的高性能 GPU kernel library。本章分析其中一条运行在 Blackwell 上、用于 prefill 的 sparse-attention forward kernel。模型做 prefill，也就是并行处理 prompt 中的 tokens 时，会先由外部 indexer 为每个 query token 选出需要关注的历史位置，再由 kernel 读取这些位置对应的 KV rows，依次完成 QK、softmax 和 PV。与 dense attention 相比，它只减少了参与计算的历史位置，attention 的主流程没有改变。 + +在本章分析的 kernel 中，一个 query token 对应 128 个 query heads，而每个被选中的历史 token 在 KV cache 中只有一条供这些 heads 共享的 row。但 128 个 heads 读取同一条 row，为什么仍能得到不同的结果？共享的是 cache 中压缩后的状态，各 head 特有的变换仍保留在 query 和 output 路径中。要看清这些变换怎样移动，先从普通 MHA 的 KV cache 算起。 + +## 从普通 MHA 的 KV cache 到 MLA + +### 普通 MHA 的 KV cache 开销 + +普通 MHA 有 $n_h$ 个 heads，每个 head 的宽度为 $d_h$，并拥有自己的 query、key 和 value projections。下面固定其中一个 head，把它的投影矩阵记为 $W^Q$、$W^K$ 和 $W^V$。在某一层中，把送入 attention projections 的当前 token 向量记为 $h_t$，把第 $s$ 个历史 token 的输入向量记为 $h_s$，则 + +$$ +q=W^Q h_t,\qquad +k_s=W^K h_s,\qquad +v_s=W^V h_s. +$$ + +这个 head 的 attention 为 + +$$ +p_s=\operatorname{softmax}_s +\left(\frac{q^{\mathsf T}k_s}{\sqrt{d_h}}\right), +\qquad +o=\sum_s p_s v_s. +$$ + +自回归生成时，当前的 query 用完即可丢弃，历史 $k_s$ 和 $v_s$ 则会被后续 query 反复读取。因此，每层需要为每个 token 缓存 $2n_hd_h$ 个元素：每个 head 各一份 K 和 V。 + +总容量会随 batch size、context length 和层数线性增长。例如，对于一个 32 层的模型，若 batch size 为 1、context length 为 4096，共有 32 个 head 且每个 head 为 128 维，那么使用 BF16 时，KV cache 就需要约 2 GB。生成每个新 token 还要重新读取这段不断增长的历史，所以 KV cache 同时带来显存容量和内存带宽压力。 + +降低这项开销的一个直接办法，是让多个 query heads 共享更多状态。Multi-query attention（MQA）让所有 query heads 共享一组 K/V，grouped-query attention（GQA）则在组内共享。MLA 采用另一种共享方式：各 head 仍使用自己的 projection，cache 中却只保存一份共享的低维状态。 + +### 从共享状态到低维 latent cache + +先看不含位置信息的 content 部分，并把当前 query 的 content 记为 $q^C$。普通 MHA 会先生成 $k_s^C=W^Kh_s$，再计算 QK 的 content score。矩阵结合律允许我们把这一步改写为 + +$$ +\mathrm{QK}_s +=(q^{C})^{\mathsf T}k_s^C +=(q^{C})^{\mathsf T}W^Kh_s +=\left((W^K)^{\mathsf T}q^{C}\right)^{\mathsf T}h_s. +$$ + +PV 同样会先生成 $v_s=W^Vh_s$，再根据 attention weights $p_s$ 做加权求和。线性变换对求和的分配性给出 + +$$ +\mathrm{PV} +=\sum_s p_s v_s +=\sum_s p_s W^Vh_s +=W^V\left(\sum_s p_s h_s\right). +$$ + +第一式改为先对当前 query 计算 $(W^K)^{\mathsf T}q^C$，再与每个 $h_s$ 做点积；第二式则先对 $h_s$ 做加权求和，再应用 $W^V$。数值结果保持不变，cache 却只需长期保存一份 $h_s$。这可以看成一个尚未压缩的 latent cache 思想实验。 + +$h_s$ 的宽度仍然是 $d_{model}$，直接缓存它虽然实现了共享，QK 打分和 PV 加权聚合却仍要在这个较宽的空间中计算。实际的 MLA 分成两步。 + +第一步，所有 heads 共用矩阵 $D$，把 $h_s$ 压缩成 $d_c$ 维的 $c_s$。就 content 部分而言，KV cache 只保存这个 $c_s$： + +$$ +c_s=Dh_s,\qquad c_s\in\mathbb{R}^{d_c}. +$$ + +第二步，每个 head 都有自己的 up-projection 矩阵。继续固定其中一个 head，把它对应的矩阵块记为 $U_K$ 和 $U_V$。如果显式展开 K/V，这个 head 的 content key 和 value 为 + +$$ +k_s^C=U_Kc_s,\qquad +v_s=U_Vc_s. +$$ + +这两个式子规定了当前 head 应该使用的 K/V。后面的 weight absorption 会保持相同的计算结果，但不再真正生成这些中间向量。$D$、$U_K$ 和 $U_V$ 会在模型训练时一起学习。由于 $d_c$ 远小于 $d_{model}$，每个 token 需要缓存的 content 状态也随之缩小。 + +### 为 RoPE 单独保留一组 Q/K 特征 + +对于不含位置信息的 content score，可以先对当前 query 计算一次 $U_K^{\mathsf T}q_t^C$，再把结果用于所有历史位置。加入 RoPE 后，首先要确认这项计算还能不能复用。 + +记 $q_t^C$ 为当前 query 的 content 特征，$R_t$ 和 $R_s$ 分别为 query 位置 $t$ 和历史位置 $s$ 的 RoPE 旋转。如果直接对 content query 和 content key 使用 RoPE，则 + +$$ +\left(R_tq_t^{C}\right)^{\mathsf T} +\left(R_sU_Kc_s\right) += +\left(U_K^{\mathsf T}R_s^{\mathsf T}R_tq_t^{C}\right)^{\mathsf T}c_s. +$$ + +当前 query 的位置 $t$ 固定，但比较不同历史 token 时，key 的位置 $s$ 会改变。上式中的 $U_K^{\mathsf T}R_s^{\mathsf T}R_tq_t^C$ 因而也随 $s$ 改变。也就是说，RoPE 一旦直接作用于 content Q/K，原本想要复用的 query-side 结果就会随历史位置变化，无法用于所有 cached rows。 + +MLA 因此另外生成一组较窄的 Q/K 特征，并只对这组特征使用 RoPE。下面用 $q_t^R$ 和 $k_s^R$ 表示它们经过位置 $t$ 和 $s$ 的 RoPE 旋转后的结果。上标 $R$ 表示这组特征经过了 RoPE；它们仍然由 token 的表示投影得到，并不是只包含位置的向量。$q_t^R$ 属于当前 head，$k_s^R$ 则由所有 heads 共享。 + +在展开 K/V 的写法中，完整 query 和 key 分别是 $[q_t^C;q_t^R]$ 与 $[U_Kc_s;k_s^R]$，分号表示沿 feature 维拼接。于是完整的 QK 点积为 + +$$ +\mathrm{score}_s +=[q_t^C;q_t^R]^{\mathsf T}[U_Kc_s;k_s^R] +=(q_t^C)^{\mathsf T}U_Kc_s + +(q_t^R)^{\mathsf T}k_s^R. +$$ + +相加的是同一个 QK 点积在两组坐标上的两个标量贡献，并不是把两组向量相加。第一项衡量 content 的相关性；第二项是经过 RoPE 的 Q/K 特征匹配，为 score 加入相对位置信息。 + +### Weight absorption：用结合律省去各 head 的 K/V 展开 + +如果直接按上面的公式计算，缓存中的每个 $c_s$ 都要先乘以 $U_K$ 和 $U_V$，显式生成当前 head 的 content K/V。Weight absorption 利用结合律改变乘法的分组，从而省去这些历史 K/V 中间结果。 + +先看上一节 score 的 content 项。Key 一侧可以改写为 + +$$ +(q_t^C)^{\mathsf T}(U_Kc_s) +=(U_K^{\mathsf T}q_t^C)^{\mathsf T}c_s. +$$ + +左边先为每个历史位置生成 content key $U_Kc_s$；右边先为当前 query 计算一次 $U_K^{\mathsf T}q_t^C$，再把结果与所有缓存中的 $c_s$ 做点积。将右式替换回上一节 score 的第一项后，经过 RoPE 的第二项保持不变，完整 score 的数值也不变。使用相同的 scale、mask 和 softmax，得到的 attention weights $p_s$ 因而相同，同时也无需把 $U_Kc_s$ 物化成各 head 的 content key。 + +Value 一侧也可以改变计算顺序。记 $p_s$ 为当前 head 对历史位置 $s$ 的 attention weight，则 + +$$ +o +=\sum_s p_s v_s +=\sum_s p_s U_Vc_s +=U_V\left(\sum_s p_s c_s\right). +$$ + +原来的顺序是先为每个历史位置生成 $v_s=U_Vc_s$，再对这些 values 加权求和。重排后，attention 先在共享的 latent 空间中计算 $\sum_s p_s c_s$，最后只对结果应用一次 $U_V$。Multi-head attention 最后还会用输出投影 $W_O$ 合并各 head 的结果；由于 $U_V$ 也是线性变换，它可以预先与 $W_O$ 中当前 head 对应的矩阵块组合。 + +Key 和 Value 两侧都只改变了乘法的结合顺序，模型参数和数值结果保持不变。这两处重排合称为 weight absorption。同一个 MLA layer 因而有两种等价的求值方式：显式生成各 head 的 K/V 称为 MHA mode；使用上述重排、直接在共享 latent 上计算称为 MQA mode。 + +```{figure} ../../img/flashmla_mla_modes_zh.svg +:width: 100% +:alt: MLA 的 MHA 与 MQA 执行模式，以及 Key 和 Value 两侧的 weight absorption + +*MHA mode 显式展开各 head 的 K/V。MQA mode 在 QK 前计算 $U_K^{\mathsf T}q_t^C$，在 latent 加权求和后应用 $U_V$；两种方式使用相同的 RoPE 项并产生相同结果。* +``` + +两种方式的结果相同，执行成本却不同。MHA mode 需要先生成各 head 的 K/V，但 QK 点积和 value 聚合处理的特征维度较小；如果许多 queries 会复用同一批展开后的 K/V，这项成本就可以被摊薄。MQA mode 不生成这些中间 K/V，却需要直接在维度较大的 latent 空间中完成 QK 和 PV。哪一种更合适，取决于当前处于 prefill 还是 decode、attention 的稀疏程度、张量形状、数据搬运成本和硬件调度方式。 + +本章研究的 sparse-prefill kernel 采用 MQA mode。 + +### MLA cache 的组成与大小 + +在 MQA mode 中，每个历史 token 的 cache 只保存共享的 $c_s$ 和经过 RoPE 的共享 key 特征 $k_s^R$，不保存各 head 展开后的 content K/V。各 head 仍使用自己的 query、$U_K$ 和 $U_V$，因此会得到不同的 attention weights 和 output。下图把普通 MHA cache 与 MLA shared cache 放在一起： + +```{figure} ../../img/flashmla_cache_story_zh.svg +:width: 100% +:alt: 普通 MHA 为每个 head 分别缓存 key 和 value；MLA 只保存一份共享压缩状态，并把各 head 特有的计算放在 attention 两侧 + +*普通 MHA 为每个 head 分别保存一份 key/value。MLA 为每个 token 只保存共享的压缩 content 状态和经过 RoPE 的 key 特征；各 head 特有的 query 与 output 变换分别在 attention 前后完成。* +``` + +现在可以比较不同机制实际缓存的内容。下面 $n_{kv}$ 表示 GQA 的 KV-head 数，$d_h^R$ 表示 $k_s^R$ 的宽度。表中的“普通 MQA”指一种 attention 结构，不是上文 MLA 的 MQA 执行方式。元素数均按每层、每个 token 计算，不计数据类型和内存对齐等额外开销： + +| 机制 | Cache 元素数 | 缓存的内容 | +| --- | ---: | --- | +| MHA | $2n_hd_h$ | 每个 head 的完整 K 和 V | +| GQA | $2n_{kv}d_h$ | 每个 KV head 的完整 K 和 V | +| 普通 MQA | $2d_h$ | 所有 query heads 共用一组完整 K/V | +| MLA | $d_c+d_h^R$ | 所有 heads 共用 $c_s$ 和 $k_s^R$ | + +MLA 的缓存大小不是 $2d_c$，因为同一个 $c_s$ 已经包含生成 content K 和 V 所需的信息，只需缓存一次。 + +本章使用 $d_c=512$、$d_h^R=64$，因此每条 $[c_s;k_s^R]$ 包含 $512+64=576$ 个标量元素。后文 kernel 中的 `d_qk=576` 指的就是这条 cached row 的宽度。 + +### 两种执行模式的数值验证 + +这个 CPU 程序同时构造两种执行方式。MHA 路径显式展开 K/V，MQA 路径吸收相同的矩阵，并且两边都加入共享的 RoPE score term。使用 Float64 可以让等价性检查足够敏感，从而发现 index 转置或 contraction 写错之类的问题。 + +```python +import math +import torch + +torch.manual_seed(0) +Q, K, H = 3, 5, 4 +D_CONTENT, D_LATENT, D_VALUE, D_ROPE, D_MODEL = 7, 6, 8, 3, 11 + +# 每个 head 的 query、每个 key token 一份共享 latent KV，以及共享 RoPE key。 +q_content = torch.randn(Q, H, D_CONTENT, dtype=torch.float64) +q_rope = torch.randn(Q, H, D_ROPE, dtype=torch.float64) +c_kv = torch.randn(K, D_LATENT, dtype=torch.float64) +k_rope = torch.randn(K, D_ROPE, dtype=torch.float64) + +W_UK = torch.randn(H, D_CONTENT, D_LATENT, dtype=torch.float64) +W_UV = torch.randn(H, D_VALUE, D_LATENT, dtype=torch.float64) +W_O = torch.randn(D_MODEL, H, D_VALUE, dtype=torch.float64) + +# MHA mode：为每个 head 显式展开 key 和 value。 +k_content = torch.einsum("hdc,kc->khd", W_UK, c_kv) +v_content = torch.einsum("hvc,kc->khv", W_UV, c_kv) +scale = 1.0 / math.sqrt(D_CONTENT + D_ROPE) +scores_mha = ( + torch.einsum("qhd,khd->qhk", q_content, k_content) + + torch.einsum("qhr,kr->qhk", q_rope, k_rope) +) * scale +prob = torch.softmax(scores_mha, dim=-1) +head_out_mha = torch.einsum("qhk,khv->qhv", prob, v_content) +model_out_mha = torch.einsum("mhv,qhv->qm", W_O, head_out_mha) + +# MQA mode：把 W_UK 移到 Q，attention 直接读取 c_kv，再把 W_UV 移到 output。 +q_absorbed = torch.einsum("qhd,hdc->qhc", q_content, W_UK) +scores_mqa = ( + torch.einsum("qhc,kc->qhk", q_absorbed, c_kv) + + torch.einsum("qhr,kr->qhk", q_rope, k_rope) +) * scale +latent_out = torch.einsum("qhk,kc->qhc", torch.softmax(scores_mqa, -1), c_kv) +W_O_absorbed = torch.einsum("mhv,hvc->mhc", W_O, W_UV) +model_out_mqa = torch.einsum("mhc,qhc->qm", W_O_absorbed, latent_out) + +torch.testing.assert_close(scores_mha, scores_mqa, rtol=1e-12, atol=1e-12) +torch.testing.assert_close(model_out_mha, model_out_mqa, rtol=1e-12, atol=1e-12) +print("weight absorption: exact up to float64 roundoff") +``` + +这里的 scale 由模型语义中的 QK head dimension 决定，因此 absorbed path 沿用原有 scale。吸收权重后的 dot product 虽然有 $D_{latent}$ 个 coordinates，$1/\sqrt{D_{latent}}$ 对应的却是另一套 scale 语义。 + +:::{admonition} Query 侧低秩分解的作用范围 +:class: note + +Query projection 也可以通过独立的低秩 latent 分解。记 down-projection 和 up-projection 为 $D_Q$ 和 $U_Q$，则 + +$$ +c^Q=D_Qh,\qquad q^C=U_Qc^Q. +$$ + +这项分解主要降低训练时的 activation memory，KV cache 大小仍由 KV-side representation 决定。它发生在 core attention 之前：attention 实现接收到的 `q` 已经是 projection 后的 query。因此，分析 KV 路径时可以直接把 $q^C$ 作为输入。 +::: + +MLA 的完整结构、训练方式和原始符号可参见 [DeepSeek-V2 论文](https://arxiv.org/abs/2405.04434)。这里保留了理解后续 FlashMLA kernel 所需的 KV compression、decoupled RoPE 和 weight absorption。 + +Weight absorption 解释了一条共享 KV row 怎样服务多个 query heads。Sparse prefill 还涉及另一项独立边界：token selection 与 sparse core attention 的职责划分。 + +## Sparse-prefill operator 的 token 选择边界 + +Dense attention 会访问每个合法的 KV token。DSA 先用轻量的 *lightning indexer* 为候选 token 打分，并为每个 query 选出一个 top-$k$ 集合；sparse core attention 随后只读取这些 latent KV entries。若原始 context length 为 $L$，prefill 的 core-attention 计算量会从 $O(L^2)$ 变成 $O(Lk)$，不过 indexer 自身也有开销。 + +```{figure} ../../img/flashmla_sparse_story_zh.svg +:width: 100% +:alt: Lightning indexer 先选择 token，sparse-prefill operator 再执行 attention + +Token 选择与 sparse core attention 是两个相互独立的算子。Indexer 输出 row addresses；sparse-prefill operator 根据这些地址 gather 对应的 rows，再完成 QK--softmax--PV 计算。 +``` + +本章研究的 sparse-prefill operator 直接接收外部 indexer 生成的 `indices` tensor，并据此执行以下语义： + +1. gather 指定的 KV rows； +2. 将越界位置和被 length mask 的位置标为 invalid； +3. 对剩余 rows 计算 attention；如果 caller 给出了重复 indices，重复项也会参与计算； +4. 返回 output、maximum logit 和 log-sum-exp。 + +Causal legality 编码在 index list 中：caller 通过只写入允许访问的 keys 来实现 causal attention。因此，稀疏选择与 causal constraint 是两项独立约束。 + +如果提供 `topk_length`，FlashMLA sparse-prefill contract 要求每个 query 的值都满足 `0 <= topk_length[q] <= topk`。这是 caller 应保证的前置条件。 + +这个 prefill 接口以不含 batch 维度的 flattened queries 为输入。每个 query token 提供一份 selected-token list，它的 `h_q` 个 query heads 共用这份 list；serving system 在调用前负责 flatten batch 或完成等价的 batch mapping。 + +CPU reference 会把这些规则固定为可独立执行的数值语义，为分析 FlashMLA 实现提供基准。 + +## 可执行的 sparse-attention reference + +CPU reference 使用以下通用 contract 和 shape 符号： + +| 符号 | 含义 | +| --- | --- | +| `s_q` | query rows 的数量 | +| `s_kv` | 可寻址 KV rows 的数量 | +| `h_q` | 每个 query row 的 query-head 数 | +| `h_kv` | 每个 KV row 的 KV-head 数；本接口要求为 1 | +| `d_qk` | QK 使用的 query/key 宽度 | +| `d_v` | value 与 output 的宽度，且 `d_v <= d_qk` | +| `topk` | 每个 query row 提供的 index slots 数量 | + +对应的 tensors 为 `q[s_q,h_q,d_qk]`、`kv[s_kv,h_kv,d_qk]`、`indices[s_q,1,topk]` 和 `out[s_q,h_q,d_v]`。可选 sink 的 shape 是 `[h_q]`，可选 `topk_length` 的 shape 是 `[s_q]`，返回的 `max_logits` 和 `lse` 都是 `[s_q,h_q]`。Regular head-128 specialization 会进一步固定 `h_q=128` 和 `d_v=512`；`h_kv=1` 已经是通用 sparse-prefill contract 的一部分。 + +在吸收权重后的 MQA contract 中，`kv[:, 0, :]` 同时提供 K 和 V：全部 `d_qk` 个坐标参与 QK，前 `d_v` 个坐标则作为 latent value。可执行 reference 将据此定义“应该算什么”。 + +还有两个边界语义会直接影响代码。第一，**attention sink** 可以看成额外加入一个 logit，但它对应的 value vector 为 0。固定一个 query 和 head，令 $x_j$ 表示第 $j$ 个 KV row 的普通 logit，$v_j$ 表示对应的 value，$a$ 表示 sink logit，$m$ 表示普通 logits 的最大值。Sink 只进入 output 的 denominator： + +$$ +O=\frac{\sum_j e^{x_j-m}v_j} +{\sum_j e^{x_j-m}+e^{a-m}}. +$$ + +Sink 不参与返回的 `max_logits` 或 `lse`。第二，如果一个 query 的 selected rows 全部无效，reference 约定 `output=0`、`max_logits=-inf`、`lse=+inf`。显式写出这项约定，可以避免在 softmax 中计算 `(-inf)-(-inf)`。 + +CPU 程序按四步实现这个 contract： + +1. 将 `indices` clamp 成可安全读取的地址，再 gather 对应 KV rows； +2. 合并地址边界与 `topk_length`，得到 validity mask； +3. 计算 QK、mask、softmax 与 PV，并把 sink 加入最终 denominator； +4. 返回 output、maximum logit 和不含 sink 的 log-sum-exp。 + +第一步为 gather 提供安全地址，第二步的 mask 才决定 row 是否参与计算。越界地址对应的 V row 还要在 PV 前清零，因为按照 IEEE arithmetic，softmax weight 为 0 时，`0 * NaN` 仍会产生 NaN。完整代码如下： + +```python +import math +import torch + +def sparse_prefill_reference( + q, kv, indices, sm_scale, d_v, *, attn_sink=None, topk_length=None +): + """Reference for q[SQ,H,D], kv[SKV,1,D], indices[SQ,1,TOPK].""" + s_q, h_q, d_qk = q.shape + s_kv, h_kv, kv_width = kv.shape + assert h_kv == 1 and kv_width == d_qk and d_v <= d_qk + assert indices.shape[:2] == (s_q, 1) + + idx = indices[:, 0].to(torch.long) # [SQ, TOPK] + topk = idx.shape[1] + assert topk > 0 + in_range = (idx >= 0) & (idx < s_kv) + safe_idx = idx.clamp(0, s_kv - 1) + focused_kv = kv[:, 0].float()[safe_idx] # [SQ, TOPK, D] + + # OOB index 只把 clamp 后的边界 row 当作安全的 gather 地址。PV 前清零这些 + # V rows，避免 sentinel row 的 NaN 经 0 * NaN 泄漏。刻意保留 topk_length + # 之后、但地址合法的 rows，以匹配 kernel 已说明的异常 NaN 行为。 + focused_v = torch.where( + in_range[:, :, None], + focused_kv[:, :, :d_v], + torch.zeros_like(focused_kv[:, :, :d_v]), + ) + + position = torch.arange(topk, device=q.device)[None, :] + if topk_length is not None: + assert topk_length.shape == (s_q,) + assert bool(((0 <= topk_length) & (topk_length <= topk)).all()) + length = ( + topk_length.to(torch.long)[:, None] + if topk_length is not None + else torch.full((s_q, 1), topk, device=q.device) + ) + valid = in_range & (position < length) + + logits = torch.einsum("qhd,qkd->qhk", q.float(), focused_kv) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], -torch.inf) + max_logits = logits.amax(dim=-1) + have_valid = valid.any(dim=-1)[:, None] + + # 避免所有 selected indices 都 invalid 时出现 (-inf)-(-inf)。 + softmax_origin = torch.where(have_valid, max_logits, torch.zeros_like(max_logits)) + weight = torch.exp(logits - softmax_origin[:, :, None]) + weight = torch.where(valid[:, None, :], weight, torch.zeros_like(weight)) + denominator = weight.sum(dim=-1) + numerator = torch.einsum("qhk,qkv->qhv", weight, focused_v) + + if attn_sink is None: + sink_term = torch.zeros_like(denominator) + else: + sink_term = torch.exp(attn_sink.float()[None, :] - softmax_origin) + out = numerator / (denominator + sink_term).clamp_min(torch.finfo(torch.float32).tiny)[ + :, :, None + ] + out = torch.where(have_valid[:, :, None], out, torch.zeros_like(out)) + + # FlashMLA 报告的 LSE 不包含 attention sink。全 invalid 时约定为 + # max_logits=-inf、lse=+inf、output=0。 + lse = torch.where( + have_valid, + max_logits + torch.log(denominator), + torch.full_like(max_logits, torch.inf), + ) + return out, max_logits, lse + + +torch.manual_seed(1) +q = torch.randn(2, 3, 6) +kv = torch.randn(9, 1, 6) +indices = torch.tensor([[[0, 3, -1, 12]], [[8, 1, 4, 2]]], dtype=torch.int32) +topk_length = torch.tensor([3, 2], dtype=torch.int32) +attn_sink = torch.randn(3) +out, max_logits, lse = sparse_prefill_reference( + q, kv, indices, 1 / math.sqrt(6), 4, + attn_sink=attn_sink, + topk_length=topk_length, +) +assert out.shape == (2, 3, 4) +assert max_logits.shape == lse.shape == (2, 3) +assert torch.isfinite(out).all() + +# OOB sentinel 不能继承其安全地址所指 row 中的 NaN。 +nan_q = torch.ones(1, 1, 2) +nan_kv = torch.tensor([[[torch.nan, torch.nan]], [[2.0, 3.0]]]) +nan_indices = torch.tensor([[[-1, 1]]], dtype=torch.int32) +nan_out, _, _ = sparse_prefill_reference(nan_q, nan_kv, nan_indices, 1.0, 2) +torch.testing.assert_close(nan_out, torch.tensor([[[2.0, 3.0]]])) +print(out.shape, max_logits.shape, lse.shape) +``` + +这段 reference 明确了 FlashMLA sparse-prefill operator 对合法输入和返回值的数值 contract。实现可以自由选择 tile 切分、storage reuse 和 pipeline overlap，同时需要在自身支持的 QK scale 范围内复现这套数值语义。 + +## Sparse prefill 在 FlashMLA 算子族中的位置 + +[FlashMLA 官方仓库](https://github.com/deepseek-ai/FlashMLA)按 selection 和 sequence stage 覆盖四类组合： + +| Selection | Sequence stage | 代表性用途 | +| --- | --- | --- | +| dense | prefill | MHA forward 与 backward | +| dense | decode | 为新生成的 queries 读取 MLA KV cache | +| token-sparse | prefill | 对 selected-token list 执行 DSA core attention | +| token-sparse | decode | 对 selected FP8 KV cache 执行 DSA inference | + +FlashMLA 是由表中四类算子组成的 library。本章先聚焦其中的 sparse-prefill operator，再深入它的 regular head-128 specialization；这条路径把算子语义直接连接到核心实现任务——将不规则 row addresses 整理成规则的 tensor-core tiles。 + +FlashMLA sparse-prefill public call 在概念上是： + +```text +out, max_logits, lse = flash_mla_sparse_fwd( + q, kv, indices, sm_scale, + d_v=512, + attn_sink=attn_sink, # optional [h_q], float32 + topk_length=topk_length, # optional [s_q], int32 +) +``` + +这些参数对应前面 reference 明确的语义：`h_q` 是 query-head 数，`s_q` 是 query-row 数，`sm_scale` 用来缩放 QK scores，`d_v` 决定 value 和 output 的宽度，`attn_sink` 可以为每个 query head 加入一个 value 为 0 的额外 logit，`topk_length` 则限定每个 query 中有效 indices 的前缀长度。这个调用返回归一化后的 output、经过 scale 的最大 logit，以及不包含 sink 的 log-sum-exp。 + +这段 public call 定义 FlashMLA sparse-prefill API，前面的 reference 将相关数值语义写成可执行形式。TIRx 是建立在 TVM 0.26 TIR 之上的 Python DSL 扩展；其中同名的 `flash_mla_sparse_fwd` 负责 registry 和 shape dispatch，按输入 shape 选择三个 SM100 phase-1 specializations 之一。它覆盖的是 public API 的 dispatch 子集。 + +Caller 负责满足这个入口的两项前置条件。第一，每个 `topk_length[q]` 都必须位于 `[0, topk]`；TIRx prefill specializations 直接使用该值，不做 clip 或逐项验证，因此大于 `topk` 的值会使实现读取到 `indices` storage 之外。第二，只因 `topk_length` 被 mask、地址本身仍合法的 row 会保留原数据，其中的 NaN 也可能进入优化路径。前面的 reference 复现了这项行为；越界地址则按 reference 的规则安全处理。 + +:::{admonition} TIRx 入口中的 `sm_scale` +:class: warning + +FlashMLA public call 在 runtime 接收 `sm_scale`；TIRx 入口则将三个 prefill specializations 的 `sm_scale` specialize 为 `1 / sqrt(d_qk)`，launch ABI 不含 scale argument。通过 `**kwargs` wrappers 传入的 `sm_scale=...` 会被静默忽略。因此，本章的 B200 示例覆盖 TIRx 入口所支持的 QK scale。需要其他 semantic QK scale 时，应将正确数值暴露为参数，或在编译时将其 specialize。 + +Weight absorption 保持该 semantic QK scale 不变。 +::: + +这三层边界将本章的实现目标限定为 TIRx dispatch 选中的一条具体 specialization。 + +## Blackwell 上的 regular head-128 案例 + +Regular head-128 案例自然衔接前面的 FlashAttention 章节：它沿用 QK--softmax--PV 主链，再加入 irregular gather、吸收权重后的 latent KV 和两个线程块之间的协作分工。 + +### Regular head-128 的 shape 与 dispatch 条件 + +Regular head-128 specialization 将共享 latent KV 的算法描述具体化为以下输入输出 shapes： + +| Tensor | Shape | Type | 含义 | +| --- | --- | --- | --- | +| `q` | `[s_q, 128, d_qk]` | BF16 | 吸收权重后的 queries | +| `kv` | `[s_kv, 1, d_qk]` | BF16 | 共享 latent/positional KV rows | +| `indices` | `[s_q, 1, topk]` | int32 | 直接 KV row indices | +| `attn_sink` | `[128]` | FP32 | 可选的 per-head sink logits | +| `topk_length` | `[s_q]` | int32 | 可选的有效 prefix length；每项位于 `[0, topk]` | +| `out` | `[s_q, 128, 512]` | BF16 | sparse-attention 结果 | +| `max_logits` | `[s_q, 128]` | FP32 | 最大 scaled logit | +| `lse` | `[s_q, 128]` | FP32 | 不含 sink 的 natural-log sum-exp | + +这里的 128 正是开篇问题中的 query-head 数，而 `kv` 中的 1 表示所有 heads 共享同一条 KV row。常见的 $d_{qk}=576$ 由 512 个 latent-content coordinates 和 64 个 RoPE coordinates 组成，`d_v=512` 则对应 latent value 宽度。后续 shape 讨论专指本章的 absorbed MQA specialization。 + +给前面定性的 mode 成本一个数值锚点：对同一 MLA layer，等价的 MHA representation 具有 $128+64=192$ 的 QK feature width 和 128 的 value/output feature width；这里的 absorbed MQA representation 则为 $512+64=576$ 和 512。若只粗略统计每个 query--key pair 在 QK 点积与 value 累加中涉及的乘加坐标数，两者是 $192+128=320$ 对 $576+512=1088$，后者约为前者的 3.4 倍。这一坐标数只作为算术宽度锚点；实际 runtime 还取决于数据移动、复用和 schedule。[DeepSeek-V3.2 report](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf) Appendix A 也说明，同一模型会根据阶段和算法采用不同 mode：DeepSeek-V3.1-Terminus 在 training 和 prefill 时使用 MHA mode，在 decode 时使用 MQA mode；DSA sparse prefill 则使用 MQA mode。 + +本节使用一组具体 shape：`s_q=1`、`s_kv=8192`、`h_q=128`、`h_kv=1`、`d_qk=576`、`d_v=512`、`topk=2048`。它表示一个 query row、128 个 query heads，以及由这些 heads 共享的 2048 个 selected-index slots。这些 slots 可能包含重复或越界地址。没有更短的 `topk_length` 时，物理调度会按 128 个 slots 一组访问 $N=16$ 个 tiles；若给出 `topk_length`，实际访问的 tile 数是 `max(ceil(topk_length / 128), 1)`。 + +每个 selected-index tile 在 kernel 中经历六步算术过程。为与源码对应，$L$ 表示原始 QK logits，$W$ 表示 BF16 未归一化指数权重，`mi` 是 online-softmax 的指数参考值，`li` 是相对于该参考值累积的 denominator，$\widetilde O$ 是尚未除以 denominator 的累计 output： + +```text +对每个包含 128 个 selected-index slots 的 tile： + 1. gather 该 tile 的 K rows + 2. gather 该 tile 的 V rows，并构造 validity mask + 3. QK：128 个 query heads × 128 个 selected-index slots → logits L + 4. 更新 mask 后的 online softmax，得到 W、参考值 mi 和分母 li + 5. 必要时 rescale running state，再累加 O~ += W @ V +所有 selected-index tiles 完成后： + 6. 用 li + sink 归一化 O~，写回 out、max_logits 和 lse +``` + +在这组具体 shape 中，前五步重复 16 轮，最后执行第六步。TIRx 实现为这六步分配不同的硬件角色。 + +### 完整源码导航 + +正文不会复制整份 device function，而是按照 QK、softmax、PV、数据搬运和同步关系展示短摘录。完整源码可以按下面的顺序阅读： + +| 阅读目标 | 源码入口 | +| --- | --- | +| 统一入口与 shape dispatch | [`flash_mla_sparse_fwd.py` lines 66--125](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/flash_mla_sparse_fwd.py#L66-L125) | +| Config、测试数据、PyTorch reference 与 launch ABI | [`sparse_prefill_head128_phase1.py` lines 66--244](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L66-L244) | +| 完整 regular head-128 device kernel | [`sparse_prefill_head128_phase1.py` lines 247--865](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L247-L865) | +| CTA-pair TMA、tcgen05 MMA 与 validity mask helpers | [`_tma.py` lines 10--60](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/_tma.py#L10-L60)、[`_gemm.py` lines 8--29](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/_gemm.py#L8-L29)、[`_mask.py` lines 10--29](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/_mask.py#L10-L29) | +| Specialize、编译、运行与数值检查 | [`sparse_prefill_head128_phase1.py` lines 868--905](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L868-L905) | + +推荐先看 dispatch 和 tensor ABI，再沿 `_kernel` 中的 WG0、WG1、WG2、WG3 分支阅读；遇到 TMA、MMA 或 mask 调用时再跳到相应 helper，最后用 `run_test` 对照输入、输出和 reference。正文中的代码块保留源码变量名和切片，并在相关段落链接到完整上下文。 + +TIRx regular head-128 实现位于 [`sparse_prefill_head128_phase1.py`](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py)。代码里的 `T` 是 TIR script namespace，`Tx` 是 GPU kernel helpers。`phase1` 沿用了对应 CUDA 实现的文件命名；在这条 regular prefill 路径中，它会由一个 kernel 直接生成完整的 `(out, max_logits, lse)`。 + +实现使用三个执行层级。一个 **CTA**（cooperative thread array）就是一个 CUDA thread block（线程块）；相邻两个 CTA 组成一个 **cluster**，可以共同发起 CTA-group tensor-core operation。一个 **warpgroup** 由 4 个 warps、共 128 个 threads 组成，并在这条 kernel 中承担一个专门角色。这里的一个 cluster 负责一个 query row。 + +在数学推导中，$p$ 表示归一化后的 softmax probability；而在源码里，`tmem_p` 和 register 变量 `p` 保存的是前面记作 $L$ 的原始 QK logits。源码中的 `s_frag` 和 `s_smem_gemm` 保存前面记作 $W$ 的未归一化指数权重。只有在 epilogue 中用 `li`（以及可选的 sink term）除 accumulated output 后，最终 output 才完成归一化。 + +较短的 TIRx 代码块用于展示 regular head-128 kernel 的局部上下文；可独立执行的 blocks 会显式说明，完整编译与数值验证集中在验证一节。 + +以下常量给出一次 tile 的尺寸、线程数和同步槽位： + +```python +B_H = 128 +B_TOPK = 128 +D_V = 512 +NUM_BUFS = 2 +NUM_THREADS = 512 +D_TQ = 384 +``` + +这些名字分别对应 kernel 的主要结构：`B_H=128` 是每个 logical tile 的 query-head 数，`B_TOPK=128` 是每次处理的 selected-index slots 数，`D_V=512` 是 output feature 宽度，`NUM_THREADS=512` 表示每个 CTA 有四个 warpgroups，`D_TQ=384` 是移入专用片上存储的 Q suffix 宽度。`NUM_BUFS=2` 为同步状态和 packed-validity mask 提供两个槽位。 + +这个 specialization 接受 512 或 576 的 `d_qk`，要求 `h_kv=1`、`d_v=512`，并要求 regular path 的 `topk` 是 128 的正整数倍。对于 128 heads，统一 front door 还会多做一次选择：`d_qk=512` 且 `topk<=1280` 时进入 small-top-k specialization；其他支持的 head-128 shapes 进入本章的 regular specialization。Head-64 shapes 使用 head-64 specialization。 + +`topk > 0` 是调用者必须满足的前置条件。统一 front door 会拒绝 `topk<=0`；各 specialization 的 `_cfg().validate()` 只检查整除性，因此直接 import 某个 specialization 时，caller 还需单独验证 `topk` 为正数。 + +Dispatch 可以在 launch GPU kernel 之前独立检查。这个代码块可直接执行，但需要按验证一节安装 `tirx-kernels`： + +```python +from tirx_kernels.flashmla.flash_mla_sparse_fwd import ( + dispatch_reason, + select_kernel, +) + +shape = dict(h_q=128, h_kv=1, d_qk=576, d_v=512, topk=2048) +assert select_kernel(**shape) == "sparse_flashmla_prefill_head128_phase1" +print(dispatch_reason(**shape)) +# sm100 h_q=128 dispatches to regular head128 phase1 +``` + +Dispatch 记录在 [`flash_mla_sparse_fwd.py`](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/flash_mla_sparse_fwd.py#L66-L120) 中；它负责选择 specialization，device schedule 则定义被选实现的执行细节。 + +### 两个 CTA 的 query-row 分工 + +难点出现在第 3、5 步的切分轴不同：QK 要形成所有 head 与 selected token 的两两组合，PV 又要沿 token 维归约，产生 512 个 value coordinates。Regular head-128 实现让两个 CTA 通过 `cta_group=2` 共同形成这块 logical tile，并在 QK 与 PV 之间旋转切分轴。 + +这个关系先用所有权图（ownership map）表示： + +```{figure} ../../img/flashmla_cta_ownership_zh.svg +:width: 100% +:alt: 两个 CTA 对 query heads、selected K rows 和 V feature columns 的 ownership + +对于每个 query row，CTA pair 会在 QK 与 PV 之间改变 logical partition。每个 CTA 最终写出 64 个完整的 output heads，每个 head 含 512 个 coordinates。 +``` + +这个 CTA pair 会用三种不同方式切分三个轴： + +| 一个 128-token top-k tile 中的资源 | CTA 0 | CTA 1 | +| --- | --- | --- | +| Query/output head ownership | heads 0--63 | heads 64--127 | +| K-row gather ownership | selected tokens 0--63 | selected tokens 64--127 | +| V-feature gather ownership | value columns 0--255 | value columns 256--511 | + +这个 2-CTA tensor-core operation 共同覆盖一块完整的 logical tile：QK 形成 128 个 heads 与 128 个 selected tokens 的两两组合，PV 再沿 token 维归约，得到 512 个 value coordinates。Collective `cta_group=2` MMA、配对的片上布局和跨 CTA 同步负责协调这两种切分。 + +源码中的 launch topology 直接实现了这张 ownership map。Launch grid 包含 `2 * s_q` 个 CTA，并将相邻 CTA 两两组成 cluster： + +```python +block_idx = T.cta_id([2 * s_q]) +T.cta_id_in_cluster([2]) +cta_idx: T.let = block_idx % 2 +s_q_idx: T.let = block_idx // 2 +thread_idx = T.thread_id([512]) +T.warpgroup_id([4]) +``` + +因此，一个 cluster 负责一个 query row，每个 CTA 含 4 个 warpgroups。这个划分还可以从后面的数据索引直接看出：Q 按 `cta_idx` chunk，K producer 选择每个 top-k block 的 `cta_idx` 半块，V producer 则从 `cta_idx * 256` 开始。 + +## Tile 的数据驻留与生命周期 + +下面的数据驻留图使用三种硬件存储：**global memory（GMEM）** 保存 kernel 的输入输出；**shared memory（SMEM）** 是 CTA 内线程共同访问的片上存储；**tensor memory（TMEM）** 是 Blackwell tensor cores 附近用于 operands 与 accumulators 的专用片上存储。 + +```{figure} ../../img/flashmla_dataflow_zh.svg +:width: 100% +:alt: QK、softmax、PV 与 epilogue 期间，global memory、shared memory、tensor memory 和 WG0 registers 中的数据驻留与生命周期复用 + +Q 被拆成 SMEM prefix 和 TMEM suffix。Gather 后的 K/V 进入 SMEM；原始 QK logits 与 output 在 TMEM 中累积；未归一化 softmax weights 再经过 SMEM 交给 PV。 +``` + +图中的 registers 归当前线程所有。**Tensor Memory Accelerator（TMA）** 负责在 GMEM 与 SMEM 之间异步搬运规则 tile，也能根据地址列表执行 gather；`tcgen05` tensor-core operation 则从 SMEM/TMEM 读取 operands，并把大型 accumulators 留在 TMEM。 + +QK 的 operand 来源用两个简写表示：**SS** 表示 Q、K 都从 SMEM 读取；**TS** 表示 Q 从 TMEM 读取、K 仍从 SMEM 读取。这条 kernel 将 Q 的 384-column suffix 搬到 TMEM，只把 prefix 留在 SMEM，所以 QK 要先做 SS prefix，再做 TS suffix。Softmax 产生的 BF16 未归一化权重则要写入 SMEM，供 PV GEMM 使用。 + +对一个 CTA 而言，重要的 logical views 如下： + +| Storage | Logical tile | Lifetime 与用途 | +| --- | --- | --- | +| SMEM `q_full` | `64 x d_qk` BF16 | Q prologue；其 prefix 留给 SS QK | +| TMEM `q_tmem` | `64 x 384` BF16 | Q 的 suffix，供 TS QK 使用 | +| SMEM `k_smem` | `64 x d_qk` BF16 | 128-row K tile 中由本 CTA gather 的一半 | +| TMEM `tmem_p` | `64 x 128` FP32 logical view | 交给 softmax 的原始 QK logits $L$ | +| SMEM `s_smem_gemm` | `64 x 128` BF16 | 交给 PV 的未归一化指数权重 $W$ | +| SMEM `v_smem_gemm` | `128 x 256` BF16 logical view | `v_smem` 的 rearranged view：所有 tile rows、本 CTA 的 V columns | +| TMEM `o_tmem` | `64 x 512` FP32 logical view | running unnormalized output | +| SMEM `o_smem` | `64 x 512` BF16 | TMA store 前的 epilogue staging | + +这些都是逻辑视图（logical views）；CTA-group TMEM layout 和 rearrangement 决定 MMA 与 load/store instructions 的实际 lane mapping。源码中的 `SMEMPool` 是共享内存分配器，用来从同一片动态 SMEM 中切出带对齐和生命周期约束的 views。源码还另行分配一个 512-column CTA-group TMEM pool，再从中切出 O、raw-logit 和 Q views。 + +SMEM 采用激进的 alias 与复用策略。只要 lifetime 允许，`q_full`、gather 后的 K/V region 和 output epilogue 就会复用类似 union 的 base。最后 384 个 Q columns 移到 TMEM 后，只有 $d_{sq}=d_{qk}-384$ 的 prefix 仍需保持 live，供 QK 第一部分使用。`d_qk=512` 时 $d_{sq}=128$；`d_qk=576` 时则为 192。具体 allocation plan 见 [`sparse_prefill_head128_phase1.py` lines 302--365](https://github.com/mlc-ai/tirx-kernels/blob/5be39749e7dfd2c4bdae9b4d396f8ec35af07126/tirx_kernels/flashmla/sparse_prefill_head128_phase1.py#L302-L365)。 + +这些存储区域能否安全地原位复用，取决于 **completion barrier**。这个小型硬件状态对象记录异步 producer 何时完成；phase bit 区分同一 barrier slot 的前后两次使用，使释放后的 K、V、$L$ 和 $W$ segments 能够安全地被原位覆盖。 + +## Warpgroup 的角色分工 + +数据放置确定后，还需要为每块 tile 指定 producer、consumer 和 storage 归还者。四个 warpgroups 由此组成 role-specialized pipeline： + +| Warpgroup | Warps | 职责 | +| --- | --- | --- | +| WG0 | 0--3 | 从 TMEM 加载原始 logits $L$，mask、online softmax、写权重 $W$、rescale O、执行 epilogue | +| WG1 | 4--7 | 加载 index fragments，并为 K 发起 gather4 TMA | +| WG2 | 8--11 | 加载 index fragments，并为 V 发起 gather4 TMA | +| WG3 | 12--15 | CTA 0 的 warp 12 发起 CTA-group QK/PV MMA；每个 CTA 的 warp 13 构造 validity mask | + +WG3 的有效工作集中在 warp 12 的 MMA issue 与 warp 13 的 validity mask。这种不对称分工与各项操作所需的并行度相匹配：一个 elected lane 为 CTA pair 发起 MMA，warp 13 负责 validity packing，WG0 则使用较多 lanes 完成 exponentiation、row reduction 和 epilogue conversion。 + +:::{admonition} Register budget 也跟着角色分配 +:class: note + +WG0 将上限提高到 144 registers，WG3 提高到 168；producer groups 则降到 96。TIRx API 用 `T.ptx.setmaxnreg(True, ...)` 表示提高，用 `T.ptx.setmaxnreg(False, ...)` 表示降低。这组配额为各 warpgroups 的角色分工提供相应的 register budget。 +::: + +### 从不规则 rows 到规则 tiles + +稀疏的 row addresses 破坏了 dense attention 所用的 contiguous 2-D copy pattern。WG1 和 WG2 使用显式 TMA `gather4`：一次 issue 提供恰好 4 个 row coordinates，让一个 warp 可以把不连续的 KV rows 搬进规则的 SMEM tile。共享 helper 固定了 CTA-pair policy。 + +Gather 使用最基本的 barrier producer--consumer handshake：producer 完成数据写入后，通过 ready/completion barrier 通知 consumer；consumer 等待后读取数据，使用完这段 storage 后再通过 done/free barrier 把复用权还给 producer。Phase ring 将这套 handshake 扩展到连续的循环迭代。 + +这段上下文摘录展示一次 `gather4` issue 读取的 addresses，以及接收其 completion 的 barrier。三个关键名字是：`cur_buf` 表示当前 tile 在两槽 barrier ring 中使用的槽位；`bar` 是 producer 与 consumer 共享的 completion barrier；`leader_mbar(...)` 取得 CTA pair 中负责汇总 TMA completion 的 leader 地址。Index names 和 slices 与链接源码一致： + +```python +_kv_gather_tma = partial( + tma_config, + dispatch="tma_explicit", + cta_group=2, + cta_mask=T.uint16(1), + cache_hint=T.uint64(0x14F0000000000000), +) + +for row_group in T.unroll(WG1_ROWS_PER_WARP): + for col_atom in T.unroll(col_count): + col = T.meta_var((col_start + col_atom) * 64) + Tx.copy_async( + k_gather_tile[ + row_group * 4 : row_group * 4 + 4, + col_atom * 64 : col_atom * 64 + 64, + ], + kv_tma[0:1, col : col + 64], + **_kv_gather_tma( + mbar=leader_mbar(bar.ptr_to([cur_buf])), + gather4=[indices_int4[row_group, lane] for lane in range(4)], + ), + ) +``` + +Gather 负责搬运 rows，validity mask 负责决定这些 rows 是否参与计算。Warp 13 的每个 active lane 加载 8 个 indices，并调用 `pack_valid_mask8`。以下两个条件同时满足时，bit $i$ 才为 1： + +$$ +0\leq\text{index}_i 0: + sq_smem = q_full.sub[:, :d_sq] + Tx.gemm_async( + tmem_p[:, :], + sq_smem[:, :d_sq], + k_smem[:, :d_sq], + **_mma_config(accum=mma_p_accumulate, smem_desc=mma_smem_desc), + ) + mma_p_accumulate = T.uint32(1) + +Tx.gemm_async( + tmem_p[:, :], + q_tmem[:, :D_TQ], + k_smem[:, d_sq : d_sq + D_TQ], + **_mma_config(accum=mma_p_accumulate, smem_desc=mma_smem_desc), +) +``` + +第一步是 SS：Q 和 K operands 都由 SMEM 描述。第二步是 TS：Q 的 384-column suffix 来自 TMEM，K 仍在 SMEM。两步写入同一个 FP32 raw-logit accumulator（源码中的 `tmem_p`）；第一步清零，第二步累加。这样拆分 Q 后，SMEM 中只需保留较小的 Q prefix，使 union allocation 成为可能，较大的 suffix 则继续走 TS path。 + +Softmax 之后，PV 是 SS GEMM：BF16 $W$ 和 V 都在 SMEM，FP32 O accumulator 则留在 TMEM。Kernel 将 V rows 和 output columns 各分成两半，四种组合共同更新全部 512 个 value coordinates。 + +## Online softmax 中的按需 O 重缩放 + +当 `topk` 大于 `B_TOPK=128` 时，一个 query row 会连续处理多个 selected-token tiles。每个 tile 都只看到一部分 scores，不能各自独立做完整 softmax；kernel 必须把前面 tiles 的状态带到下一轮。前面的具体例子在没有更短 `topk_length` 时有 $N=16$ 个 tiles，因此会递推合并 16 轮状态。 + +为了直接使用硬件 `exp2`，先把当前 tile 中的每个原始 QK dot product $x_j$ 转换到以 2 为底的指数单位： + +$$ +r_j=x_j\cdot\text{semantic\_QK\_scale}\cdot\log_2(e). +$$ + +对于连续到来的 score tiles，online softmax 会保存指数参考值 $m$、denominator $\ell$ 和未归一化 output $\widetilde O$。合并下一块 tile 时： + +$$ +m'=\max(m,\max_j r_j),\qquad +\alpha=2^{m-m'}, +$$ + +$$ +\ell'=\alpha\ell+\sum_j2^{r_j-m'},\qquad +\widetilde O'=\alpha\widetilde O+\sum_j2^{r_j-m'}v_j. +$$ + +在 TIRx regular head-128 specialization 中，编译期 `sm_scale_div_log2` 就是 `(1 / sqrt(d_qk)) * log2(e)`。这与原来的 softmax 相同，只是改写成直接映射到快速 base-2 exponential instructions 的形式。 + +对应到源码，`cur_pi_max` 是当前 128-token tile 的 base-2 最大值，`mi` 是递推中实际采用的指数参考值，`li` 是相对于 `mi` 累积的 denominator。`real_mi` 则单独保存迄今为止真实的最大值，用于最终报告 `max_logits`。区分 `mi` 与 `real_mi`，使 lazy rescaling 能够保持输出统计量不变。 + +每次 row maximum 增加时都重缩放完整的 512-coordinate O tile，代价会很高。Head-128 kernel 因此使用 lazy threshold： + +```python +should_scale_o: T.bool = ( + T.ptx.any_sync(T.uint32(0xFFFFFFFF), cur_pi_max - mi > 6.0) != 0 +) + +if not should_scale_o: + scale_for_old = 1.0 + new_max = mi +else: + new_max = T.max(cur_pi_max, mi) + scale_for_old = T.ptx.exp2(mi - new_max) +``` + +如果当前 tile 的最大值比保存的指数参考值至多高 6 个 base-2 units，kernel 会继续保留旧参考值。新的 exponential 此时最大可能达到 $2^6=64$，已累积的 O 保持原 scale。一旦差值超过 6，kernel 就更新参考值，同时 rescale $\ell$ 和已经存在的 O。Warp-wide `any_sync` 让参与计算的 rows 使用一致的决策。 + +`real_mi` 始终维护迄今为止真实的最大值，所以 `max_logits` 保持原有语义。结束时，两个 64-token half 对每个 logical row 的贡献会被合并，kernel 输出： + +$$ +\mathrm{lse}=m\ln 2+\ln\ell. +$$ + +可选的 attention sink 将最终 output scale 改成： + +```python +output_scale: T.float32 = T.cuda.fdividef( + T.float32(1.0), li + T.ptx.exp2(attn_sink_log2 - mi) +) +``` + +这项 sink 修正只作用于最终 output scale，报告的 LSE 保持不变。对于全 invalid row，特殊分支会输出 0，同时令 `max_logits=-inf`、`lse=+inf`，与前面的 reference 一致。 + +## Pipeline 各阶段的安全重叠 + +这个 schedule 的主线是数据所有权在各角色之间的交接。每个 tile 依次经历 `K ready → QK done → L consumed → W ready → PV done → V/O reusable`：QK 完成后 softmax 才能消费 logits，softmax 产生 weights 后 PV 才能开始。这些依赖仍然留下了重叠空间；K 或 V segment 一旦释放，负责 gather 的 warpgroups 就可以继续推进。 + +源码用 **memory barrier（mbarrier）** 表示这些交接。Producer 在数据或异步操作完成时贡献 arrival，consumer 等待相应 phase；consumer 用完后，再通过 done/free barrier 把覆盖这段 storage 的权利交还 producer。Ready barrier 授予读取权，done/free barrier 授予下一轮覆盖权。 + +```{figure} ../../img/flashmla_pipeline_stages_zh.svg +:width: 100% +:alt: 相差一个 tile 的 pipeline 填充、稳态与排空，以及 QK、softmax、PV 的重叠 + +填充阶段发起 QK(0)；排空阶段发起 PV($N-1$)，随后进入最终 epilogue。稳态中，softmax($k-1$) 可与 QK($k$) 重叠；唯一的 MMA issuer 串行发出 QK($k$) 与 PV($k-1$)；QK($k$) 完成后，softmax($k$) 可与仍在异步执行的 PV($k-1$) 重叠。Tensor-core issue stream 因而只有一条，重叠来自异步执行和其他 warpgroups 的并行工作；前面的具体例子中 $N=16$。 +``` + +这张宏观时间线由初始化和一组具体 barrier edges 实现。Kernel 先由 warp 0 初始化 mbarriers，执行 cluster sync，launch Q prologue，分配 CTA-group TMEM，再进入 specialized loops。 + +CTA-group gather 的 TMA completion 会被路由到指定的 leader barrier，使 CTA pair 发出的操作共同满足同一 expected byte count。 + +主要 barrier edges 对应以下 storage ownership transfers： + +| Barrier | Producer 到 consumer | 保护的 storage | +| --- | --- | --- | +| `bar_k_part0_ready` | WG1 到 WG3 | SS QK 使用的 K prefix | +| `bar_qk_part_done` | WG3 到 WG1 | SS QK 完成后允许覆盖 K prefix | +| `bar_k_part1_ready` | WG1 到 WG3 | TS QK 使用的 K suffix | +| `bar_qk_done` | WG3 到 WG0 和 WG1 | 原始 logits $L$ ready；QK 完成后 K suffix 可复用 | +| `bar_p_free` | WG0 到 WG3 | 下一次覆盖前，TMEM raw-logit tile 已被消费 | +| `bar_k_valid_ready/free` | warp 13 到/从 WG0 | packed validity mask | +| `bar_so_ready` | WG0 到 WG3 | PV 所需的 BF16 权重 $W$ ready | +| `bar_v_part0_ready` / `bar_sv_part_done` | WG2 到/从 WG3 | 第一半 V | +| `bar_v_part1_ready` / `bar_sv_done` | WG2 到 WG3，再由 WG3 到 WG2 和 WG0 | 第二半 V；复用 V、rescale O 或执行 epilogue 前，PV/O 已完成 | + +Barrier 槽位按 tile 编号循环使用： + +```python +cur_buf = k % 2 +cur_phase = (k // 2) & 1 +``` + +这样，复用的 barrier slot 可以区分本轮到达与两轮之前的旧到达。 + +`bar_qk_part_done` 允许 producer 在 K suffix 可以复用之前先替换 K prefix。两条 `bar_sv_*` edge 对 V 做同样的事。 + +```{figure} ../../img/flashmla_pipeline_zh.svg +:width: 100% +:alt: Sparse-prefill 详细 pipeline，展示 QK 与 PV 的串行发起、K/V 分段复用、mask-slot ring 和 WG0 交接 + +在稳态中，QK($k$) 与 PV($k-1$) 交错发起，其他 roles 同时执行分段 K/V gather 和 softmax。Barrier phase 保护单份 tile 的 in-place storage reuse。 +``` + +同一片 SMEM 在不同 memory proxy 之间交接时，正确可见性同时依赖 completion signaling 和 proxy ordering。线程执行的普通 SMEM load/store 属于 **generic proxy**，TMA 与 tcgen05 的异步访问属于 **async proxy**。 + +因此，`T.ptx.tcgen05.fence.*` 约束 TMEM access 与线程可见操作的先后关系；`T.ptx.fence.proxy_async("shared::cta")` 则在 SMEM 的 generic 与 async proxy 之间建立顺序。它既用于普通 store 写入 $W$ 或 epilogue tile 后、tcgen05/TMA 执行异步读取之前，也用于异步 SMEM read 完成后、普通代码覆写共用 storage 之前。 + +Mbarrier 传达 completion 并移交 storage 使用权，proxy fence 约束不同 proxy 的 memory effects；两者共同完成一次安全交接。 + +## Regular head-128 的编译与数值验证 + +完整验证覆盖三层：regular head-128 实现能够用 TVM 0.26 编译，生成的 kernel 能够在 B200 上实际 launch，并且它的 output、maximum logits 和 LSE 与前面的 reference 一致。 + +这个 specialization 面向 compute capability 10，其 TMA/tcgen05 形式要求 SM100 class GPU。环境应使用 B200、CUDA 12.9 或更高版本，以及官方 Apache TVM 0.26.0 package。 + +首先通过 [PyTorch 官方选择器](https://pytorch.org/get-started/locally/)安装支持 B200 的 CUDA-enabled PyTorch build。`tirx-kernels` 会 import PyTorch，因此需要单独安装这个未声明的 package dependency。 + +随后安装 TVM 和 kernel repository，再运行示例： + +```bash +python -m pip install "apache-tvm==0.26.0" cuda-bindings +git clone https://github.com/mlc-ai/tirx-kernels.git +cd tirx-kernels +git checkout 5be39749e7dfd2c4bdae9b4d396f8ec35af07126 +pip install -e . +``` + +只用一个 query row 的 smoke test 仍会覆盖完整的 128-head、top-k-2048 kernel，同时保持 reference 时间较短： + +```python +from tirx_kernels.flashmla.sparse_prefill_head128_phase1 import run_test + +run_test( + label="tutorial_smoke", + s_q=1, + s_kv=8192, + topk=2048, + d_qk=576, + h_q=128, + h_kv=1, + d_v=512, + have_attn_sink=True, + have_topk_length=False, + seed=0, +) +print("compile, launch, and randomized reference check passed") +``` + +`run_test` 同时覆盖这三层验证：它会分配随机 BF16 Q/KV 和随机 indices，编译并运行生成的 kernel，逐个 query row 求 FP32 PyTorch oracle，再用明确的 tolerance 检查 output、maximum logits 和 LSE。因此，它可以作为默认的 end-to-end gate，覆盖单纯 PTX 编译无法发现的 head partition 错误或遗漏 validity bit。 + +`tirx-kernels` CLI 可以运行完整的 registered configuration： + +```bash +python -m tirx_kernels.test \ + --kernel sparse_flashmla_prefill_head128_phase1 \ + --config bench_regular_dqk576_hq128_s4096_kv8192_topk2048 +``` + +Negative tests 同样重要。设置 `inject_invalid_indices=True` 可以覆盖负数和过大的 row ID，设置 `have_topk_length=True` 可以覆盖 position predicate；全 invalid row 则用于确认约定的 0/-infinity/+infinity 行为。对 head-64 和 small-top-k shapes 调用统一 dispatch entry，可以把覆盖范围扩展到其他 prefill specializations。完整 FlashMLA API 还包含 runtime `sm_scale` 等 TIRx 入口未覆盖的 call semantics，需要单独验证。 + +这些正向与异常输入测试分别验证 operator semantics 与 specialization schedule contract，二者可以归纳为五条不变量。 + +## Operator 与 specialization 的不变量 + +1. **缓存不变量（cache invariant）。** 一条 `h_kv=1` 的 latent KV row 可以服务全部 `h_q` 个 query heads，是因为 key up-projection 已吸收到各 head 的 query 路径，而 value up-projection 被移到 core attention 之后。RoPE channel 仍保持显式，QK scale 仍是模型语义规定的 scale。 + +2. **稀疏契约不变量（sparse-contract invariant）。** Token selection 发生在 sparse-prefill operator 之前；`indices` 提供 rows，重复项仍按重复项计算，因果合法性由 caller 保证，并且每个 `topk_length` 都必须位于 `[0, topk]`。优化路径与 reference path 必须保持相同的 attention sink 和全 invalid 约定。 + +3. **所有权不变量（ownership invariant）。** 在 regular head-128 specialization 中，一个 2-CTA cluster 负责一个 query row。这对 CTA 沿不同的轴切分 Q/output heads、selected K rows 和 V features；CTA 0 的 warp 12 是 CTA-group MMA 的唯一 issuer。 + +4. **驻留不变量（residency invariant）。** 在这个 specialization 中，$L$ 表示 FP32 原始 logits，$W$ 表示 BF16 未归一化指数权重。K、V、$L$ 和 $W$ 的大型 workspace 都是原位复用的单份 tile，而 `NUM_BUFS=2` 驱动的是 barrier/phase ring。在数据缓冲区（data buffers）中，只有小型 packed-validity mask 才有两个 physical slots。 + +5. **交接不变量（handoff invariant）。** 在这个 specialization 中，ready/done barriers 转移每个可复用 segment 的 ownership，proxy fences 则对 generic 与 asynchronous SMEM access 排序。唯一 issuer 保证 QK($k$) 先于 PV($k-1$) 发射，而 threshold 为 6 的优化保证上报 maximum 的语义不变，并保持 online LSE 的递推关系。 + +前两条不变量定义 FlashMLA sparse-prefill operator semantics，后三条定义 regular head-128 specialization 的 schedule contract。其他由 dispatch 选中的 specializations 可以采用各自的 tile sizes、register budgets、ownership 和 barrier topology，同时保持相同的 operator 结果。 + +## 练习与扩展验证 + +Regular head-128 specialization 是 dispatch space 中的一个点。源码树还包含 head-64 phase-1 specialization，以及 head-128 `d_qk=512` small-top-k specialization。它们采用不同 schedule，dispatch 会据此按 tile economics 选择 specialization。 + +1. **复现 weight absorption。** 给可执行的 absorption proof 加入 causal mask。确认 MHA mode 与 MQA mode 仍然一致，再故意把 absorbed path 的 scale 改为 $1/\sqrt{D_{latent}}$，测量产生的误差。 + +2. **压力测试 sparse validity。** 在 `sparse_prefill_reference` 中加入全 invalid query、重复 indices、`topk_length=0`，以及正负无穷的 sink values。分别写出预期的 `(out, max_logits, lse)`。 + +3. **追踪 ownership。** 对一个 top-k tile，给 Q、K、$L$、$W$、V 和 O 的每个 dimension 标注 `(CTA, local row, local column)`，找出一个 logical row 在哪些位置需要另一个 CTA 拥有的信息。 + +4. **审计 residency。** 从 `SMEMPool` allocation 开始画出每个 alias interval。验证为什么 $d_{sq}$ Q prefix 必须保持 live，而 384-column suffix 可以移到 TMEM，并找出结束每个 reuse hazard 的 barrier。 + +5. **测量 threshold。** 对 random 和 adversarial logits 加入 instrumentation，统计 `should_scale_o` 取 true 的频率。比较 threshold 6、always-rebase 和 never-rebase 三种策略的 numerical error 与 TMEM O traffic。 + +6. **比较 dispatch。** 对 head count 64/128、`d_qk` 512/576，以及 1280 附近的 top-k values 调用 `select_kernel`。运行前先预测选择的 specialization，再检查哪些 constraint 属于 front door，哪些由单个 specialization 强制执行。 + +7. **阅读 generated program。** 编译 smoke shape，在生成的 PTX 中找到 `tcgen05` MMA、TMA gather、mbarrier 和 proxy-fence instructions，再把它们映射回对应 TIRx line。最后重新运行 numerical check：source、generated code 和实际观察值是三类互补的证据。 + +FlashMLA 展示了一种通用的高性能 irregular-operator 设计模式：indexer 生成 sparse addresses，TMA 把这些 addresses gather 成 dense tiles，tensor cores 消费 tiles，显式 barrier 则保护激进的 storage reuse。把 algorithm、dispatch contract、ownership map 和 memory protocol 放在一起分析，便能将高速 kernel 还原成一套可以解释和验证的程序。 diff --git a/zh/index.md b/zh/index.md index 63337fd0..1c8e6553 100644 --- a/zh/index.md +++ b/zh/index.md @@ -4,7 +4,7 @@ 要让这些 kernel 真正跑得快，不能只罗列优化技巧。近年来的 GPU 架构引入了更丰富的内存空间、新的数据搬运机制和越来越专用化的执行单元。要充分利用这些硬件能力，既要理解 GPU 如何执行程序，也要掌握一个基础 kernel 如何逐步演变成高性能实现。本书将围绕这两个方面展开。 -本书按照从硬件、编程模型到完整 kernel 的顺序展开。我们会先介绍 GPU 的组织方式和执行模型，再学习本书使用的编程模型，最后逐步构建高性能 kernel。本书主要面向 NVIDIA Blackwell，并以 General Matrix-Matrix Multiplication（GEMM）和 FlashAttention 为贯穿全书的示例。在构建这些 kernel 的过程中，还会系统介绍数据布局、异步数据搬运和异步协作等关键主题。 +本书按照从硬件、编程模型到完整 kernel 的顺序展开。我们会先介绍 GPU 的组织方式和执行模型，再学习本书使用的编程模型，最后逐步构建高性能 kernel。本书主要面向 NVIDIA Blackwell，并以 General Matrix-Matrix Multiplication（GEMM）、FlashAttention 和 FlashMLA 为贯穿全书的示例。在构建这些 kernel 的过程中，还会系统介绍数据布局、异步数据搬运和异步协作等关键主题。 本书内容源自卡内基梅隆大学的 [Machine Learning Systems](https://mlsyscourse.org/) 课程系列。书中的示例使用 TIRx Python DSL，让读者能够在真实 kernel 中学习、运行和验证这些概念。TIRx 会明确表示与硬件执行有关的选择，因此可以结合可运行的代码分析控制流、内存访问和同步逻辑。 @@ -16,7 +16,7 @@ - **第一部分：理解 GPU。** 这一部分介绍 GPU 的整体架构组织、编写高性能 kernel 的通用方法，以及数据布局、异步内存操作和协作等关键概念，并建立后续章节所依赖的硬件理解。 - **第二部分：TIRx 概览。** 这一部分介绍 TIRx 的核心组成部分，为理解后续章节中的代码示例做准备。 - **第三部分：GEMM：从 Tiled 到 SOTA。** 这一部分完整讲解如何优化一个 tiled GEMM，并逐步加入 TMA pipelining、persistent scheduling、warp specialization 和 2-CTA cluster。 -- **第四部分：Flash Attention 4。** 这一部分基于第三部分的技术构建完整的 attention kernel：两个 MMA，中间插入 softmax，并包含 online-softmax rescaling、causal mask 和 GQA。 +- **第四部分：注意力内核。** 这一部分首先基于第三部分的技术构建 Flash Attention 4，随后介绍 Multi-head Latent Attention，并将一个稀疏 FlashMLA prefill 算子映射为 Blackwell 上的 2-CTA kernel。 - **附录。** TIRx 语言参考、可复现的 GPU 性能测量与分析流程、编译器内部机制，以及异步 kernel 调试指南。 ```{toctree} @@ -52,10 +52,11 @@ chapter_gemm_advanced/index ``` ```{toctree} -:caption: 第四部分：Flash Attention 4 +:caption: 第四部分：注意力内核 :maxdepth: 2 chapter_flash_attention/index +chapter_flashmla/index ``` ```{toctree}