From 35cf0a3d5e1d00912b4156fa98af3e63f9c3858a Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Wed, 5 Aug 2026 19:24:19 -0400 Subject: [PATCH] Polish and align English and Chinese documentation --- appendix/debugging_warp_specialized.md | 18 +++--- appendix/index.md | 6 +- chapter_background/index.md | 15 ++--- chapter_data_layout/index.md | 8 +-- chapter_flash_attention/index.md | 32 +++++----- chapter_gemm_advanced/index.md | 6 +- chapter_gemm_basics/index.md | 2 +- chapter_layout_generations/index.md | 2 +- chapter_performance/index.md | 10 +-- chapter_tensor_cores/index.md | 2 +- chapter_tirx_layout_api/index.md | 4 +- index.md | 49 ++++++++------- tirx_guide/arch/lowering_pipeline.rst | 16 ++--- .../language_reference/cuda/buffers.rst | 9 +-- .../language_reference/cuda/control_flow.rst | 17 +++--- .../language_reference/cuda/data_types.rst | 2 +- .../language_reference/cuda/threads_sync.rst | 18 +++--- zh/appendix/debugging_warp_specialized.md | 47 +++++++------- zh/appendix/index.md | 4 +- zh/chapter_async_barriers/index.md | 2 +- zh/chapter_background/index.md | 8 +-- zh/chapter_clc/index.md | 4 +- zh/chapter_data_layout/index.md | 4 +- zh/chapter_flash_attention/index.md | 38 ++++++------ zh/chapter_gemm_advanced/index.md | 36 +++++------ zh/chapter_gemm_async/index.md | 48 +++++++-------- zh/chapter_gemm_basics/index.md | 56 ++++++++--------- zh/chapter_intro_tirx/index.md | 2 +- zh/chapter_layout_generations/index.md | 2 +- zh/chapter_performance/index.md | 8 +-- zh/chapter_tensor_cores/index.md | 2 +- zh/chapter_tirx_layout_api/index.md | 4 +- zh/chapter_tma/index.md | 2 +- zh/chapter_tmem/index.md | 2 +- zh/index.md | 22 +++---- zh/tirx_guide/arch/lowering_pipeline.rst | 16 ++--- .../language_reference/cuda/buffers.rst | 61 ++++++++++--------- .../language_reference/cuda/control_flow.rst | 21 ++++--- .../language_reference/cuda/threads_sync.rst | 18 +++--- 39 files changed, 313 insertions(+), 310 deletions(-) diff --git a/appendix/debugging_warp_specialized.md b/appendix/debugging_warp_specialized.md index 1a646770..f86ab8e2 100644 --- a/appendix/debugging_warp_specialized.md +++ b/appendix/debugging_warp_specialized.md @@ -1,9 +1,9 @@ (chap_warp_spec_debug)= # Debugging Warp-Specialized Kernels -GEMM Steps 7-9 in {ref}`chap_gemm_advanced` overlap TMA load, `tcgen05` MMA, and TMEM/SMEM writeback. The same debugging method applies to Flash Attention handoffs: identify the roles, identify the storage each role owns, then verify the generated CUDA against that model. +GEMM Steps 7–9 in {ref}`chap_gemm_advanced` overlap TMA load, `tcgen05` MMA, and TMEM/SMEM writeback. The same debugging method applies to the handoffs among QKᵀ MMA, softmax, PV MMA, and correction in Flash Attention: identify the roles and the storage each role owns, then verify the generated CUDA against that model. -Do not start by rewriting the kernel. First make sure the run is valid, then inspect the generated CUDA. After environment and compile-time issues are ruled out, runtime failures in these kernels usually reduce to a broken handoff: an uninitialized barrier, the wrong arrival count, a collective hidden inside a role guard, a stale barrier phase, or storage reused before the producer has made its writes visible. +Do not start by rewriting the kernel. First verify the environment and reproduce the failure with the smallest correctness test, then inspect the generated CUDA. After environment and compile-time issues are ruled out, runtime failures in these kernels usually reduce to a broken handoff: an uninitialized barrier, the wrong arrival count, a collective hidden inside a role guard, a stale barrier phase, or storage reused before the producer has made its writes visible. ## Before Debugging the Kernel @@ -27,7 +27,7 @@ These kernels target Blackwell (`sm_100a`). If Python imports a stale TVM checko 7. Change one handoff at a time: init count, arrive/wait phase, role guard, fence, TMA store drain, TMEM alloc/dealloc, or tile-scheduler advance. 8. Re-run correctness before measuring performance. -## What Transfers +## Map the Data Handoffs For any asynchronous kernel, make a small worksheet before changing code: @@ -44,15 +44,15 @@ Then verify the generated CUDA against the worksheet: - Barrier inits appear before guarded role branches. - Collective operations are not accidentally narrowed by lane, warp, or warpgroup guards. - Arrive/wait phases match the handoff table. -- TMA store drains, TMEM dealloc, and SMEM reuse happen only after the lifetime table says they are legal. +- TMA store completion is awaited, TMEM is deallocated, and SMEM is reused only when the lifetime table says each action is safe. -Use the same worksheet for TMA->MMA->writeback GEMM pipelines and for the score/softmax/value/correction handoffs in Flash Attention. +Use the same worksheet for TMA → MMA → writeback GEMM pipelines and for the handoffs among QKᵀ MMA, softmax, PV MMA, and correction in Flash Attention. ## If Compilation Fails Fix compile-time failures before debugging runtime synchronization: -| Symptom | Likely area | First check | +| Symptom | Likely cause | First check | |---|---|---| | Unknown TIRx API or attribute error | Installed wheel does not match the tutorial code | Print `tvm.__file__` and `tvm.__version__`; compare the API name with {ref}`chap_language_reference`. | | Unsupported `dispatch=` | The selected target or primitive does not support that path | Check the `dispatch` argument and target capability; `tcgen05` paths in this tutorial require Blackwell. | @@ -133,7 +133,7 @@ Check these before changing the algorithm: Start from the symptom, but treat it as a clue rather than a final diagnosis: -| Clue | Likely area | First check | +| Clue | Likely cause | First check | |---|---|---| | Kernel hangs, then the runtime reports an unspecified launch failure | Deadlock | Barrier init placement, arrival count, `cta_sync()` placement, and `next_tile()` participation | | Illegal memory access, XID, or later unrelated CUDA calls also fail | Crash / poisoned context | Restart Python, then check pointer ranges, storage lifetime, and collective participation | @@ -185,13 +185,13 @@ Classify wrong output by pattern before guessing. Whole row stripes often point - **Missing `fence.proxy_async("shared::cta")` before TMA store.** The TMA engine may not see SMEM writes from threads. - **Missing `cp_async.bulk.commit_group()` plus `wait_group(0)` after TMA store.** The next tile can reuse Dsmem before the store drains. - **Persistent kernel fails intermittently at small sizes such as 1024x1024.** Larger sizes can mask the race with longer K-loops. Re-check phase reset between tiles and the TMA-store commit/wait. -- **Missing `fence.after_thread_sync()` between MMA completion and TMEM load.** The `mma2ld` wait confirms that the MMA has completed, but a writeback thread still needs `T.ptx.tcgen05.fence.after_thread_sync()` before issuing `tcgen05.ld`. This orders the new thread's TMEM load after the cross-thread completion notification. Steps 7-9 place the fence immediately after `mma2ld.wait`. This is a `tcgen05` ordering rule; it does not wait for a TMA load or make ordinary thread writes visible to the TMA engine. Those handoffs use their own mbarrier and proxy-fence protocols. +- **Missing `fence.after_thread_sync()` between MMA completion and TMEM load.** The `mma2ld` wait confirms that the MMA has completed, but a writeback thread still needs `T.ptx.tcgen05.fence.after_thread_sync()` before issuing `tcgen05.ld`. This orders the new thread's TMEM load after the cross-thread completion notification. Steps 7–9 place the fence immediately after `mma2ld.wait`. This is a `tcgen05` ordering rule; it does not wait for a TMA load or make ordinary thread writes visible to the TMA engine. Those handoffs use their own mbarrier and proxy-fence protocols. ## Correct but Slow If the output is correct but performance is far below expectation, use the same inspection loop: -| Clue | Likely area | First check | +| Clue | Likely cause | First check | |---|---|---| | Generated CUDA has no `cp.async.bulk.tensor` | Copy did not lower to TMA | Check `dispatch="tma"`, target capability, and operand layout | | Generated CUDA has no `tcgen05` path | MMA did not lower to Blackwell Tensor Core instructions | Check `dispatch="tcgen05"`, target capability, and operand layouts | diff --git a/appendix/index.md b/appendix/index.md index ca4de4f5..d836e0bd 100644 --- a/appendix/index.md +++ b/appendix/index.md @@ -1,7 +1,7 @@ (chap_appendix)= # Overview -The main path runs through Parts I–IV. The Reference holds material you reach for while reading: +The main text runs through Parts I–IV. The Reference section collects material you may want to consult while reading: | Need | Where | |------|-----| @@ -12,5 +12,5 @@ The main path runs through Parts I–IV. The Reference holds material you reach For the complete `tvm.tirx` Python API, see the [upstream TVM documentation](https://tvm.apache.org/docs/). -The TIRx native level ({ref}`chap_tirx_primer`) and the tensor layout model -({ref}`chap_tirx_layout_api`) are covered in Part II. +Part II covers the TIRx programming model ({ref}`chap_tirx_primer`) and the +tensor layout model ({ref}`chap_tirx_layout_api`). diff --git a/chapter_background/index.md b/chapter_background/index.md index 21ee8b31..fad22adc 100644 --- a/chapter_background/index.md +++ b/chapter_background/index.md @@ -97,8 +97,8 @@ A cluster can contain CTAs running on different SMs. Each CTA still owns its own This capability avoids unnecessary round trips through GMEM. One CTA can directly access another CTA's SMEM without requiring the owner to write the data back to GMEM for the peer to reload. When -an asynchronous operation moves such data, a completion barrier notifies later computation after the -transfer finishes. +an asynchronous operation moves such data, it updates a completion barrier after the transfer +finishes; consumers wait on that barrier before using the result. The figure below shows the DSMEM access path in a 2-CTA cluster. Each CTA retains its own SMEM but can read the other CTA's SMEM. @@ -115,8 +115,8 @@ In the 2-CTA GEMM shown above, each CTA stores its own slices of A and B and rea slice through DSMEM. Here, sharing does not merge the two SMEM allocations. It means only that CTAs in the same cluster can access one another's data across SMs. -The two CTAs can also form `cta_group=2` and execute a cooperative MMA that produces a larger output -tile. +The two CTAs can form a CTA pair and execute a cooperative MMA in `cta_group::2` mode, producing a +larger output tile. ## Compute: CUDA Cores and Tensor Cores @@ -139,8 +139,8 @@ and accumulator placement. Hopper introduced asynchronous warpgroup MMA (`wgmma. Blackwell's fifth-generation Tensor Core, `tcgen05`, stores accumulators in Tensor Memory rather than registers. Later chapters discuss these differences in detail. -Clusters also introduce two important GEMM uses. **2-CTA cooperative MMA** allows two CTAs to each -provide part of the SMEM operands and jointly issue a larger Tensor Core MMA tile. **TMA multicast** +Clusters enable two forms of collaboration that are important for GEMM. **2-CTA cooperative MMA** +allows two CTAs to each provide part of the SMEM operands for a larger Tensor Core MMA tile. **TMA multicast** allows one GMEM load to deliver the same tile to multiple CTAs, avoiding redundant global memory traffic from each CTA loading the same data separately. Both rely on the cluster and DSMEM mechanism introduced earlier. @@ -164,7 +164,8 @@ A single GEMM tile usually flows through three stages. 1. **Load:** A TMA copy moves an A or B operand tile from GMEM to SMEM. One thread issues the copy and records the expected number of arriving bytes. As data reaches SMEM, the TMA engine updates the - progress count. The completion barrier fires only after all expected bytes have arrived. + progress count. The completion barrier becomes complete only after all expected bytes have + arrived. 2. **Compute:** A `tcgen05` MMA reads operand tiles from SMEM and accumulates the product into a TMEM tile. One designated thread commits the MMA; when computation completes, the hardware signals the corresponding barrier. diff --git a/chapter_data_layout/index.md b/chapter_data_layout/index.md index da561f3b..40e25ad9 100644 --- a/chapter_data_layout/index.md +++ b/chapter_data_layout/index.md @@ -13,8 +13,8 @@ Computations over the same values can differ in performance by an order of magni depending only on how those values are physically arranged in memory. A tensor's logical indices do not say where its bytes are actually stored. The hardware is highly -sensitive to that placement. It determines whether loads from 32 lanes coalesce into one transaction -or split across as many as 32, whether addresses land in different memory banks or collide and +sensitive to that placement. Depending on the access pattern, it determines whether loads from 32 +lanes can be coalesced into one transaction or must be split across as many as 32, whether addresses land in different memory banks or collide and serialize, and whether a tile has a byte arrangement that a Tensor Core can read. Machine learning programs usually describe a tensor by its logical shape. A **data layout** supplies @@ -511,8 +511,8 @@ For example, row 1 in our demo contains the logical column labels arrangement, producing `9, 8, 11, 10, 13, 12, 15, 14`. We call each 128-bit cell in the figure a 16 B **sector**. In `SWIZZLE_128B`, each row of an atom -contains eight sectors, for a total width of 128 B. At the common 4-byte bank granularity, one sector -spans four banks, so a full row covers all 32 banks. The swizzle uses the row coordinate to +contains eight sectors, for a total width of 128 B. At the common 4-byte bank granularity, the four +32-bit words in one sector map to four adjacent banks, so a full row covers all 32 banks. The swizzle uses the row coordinate to XOR-permute the eight sectors within that row. A `SWIZZLE_128B` atom contains eight rows, so its total size is `8 × 128 B = 1024 B`. Here, diff --git a/chapter_flash_attention/index.md b/chapter_flash_attention/index.md index 58bf6158..d1aa2672 100644 --- a/chapter_flash_attention/index.md +++ b/chapter_flash_attention/index.md @@ -33,13 +33,13 @@ The dot product of $q_i$ and $k_j$ gives the scalar score at position $(i,j)$: $$s_{ij}=q_i\cdot k_j$$ -Fixing query vector $q_i$ and taking its dot product with every key vector $k_j$ produces the scores $s_{ij}$ for that query. These scores form row $i$ of the score matrix $S=QK^\top$. Basic online softmax first takes the largest score in that row as the reference for exponentiation: +Fixing query vector $q_i$ and taking its dot product with every key vector $k_j$ produces the scores $s_{ij}$ for that query. These scores form row $i$ of the score matrix $S=QK^\top$. Let $m_i^{\max}$ denote the exact largest score in that row: -$$m_i=\max_j s_{ij}$$ +$$m_i^{\max}=\max_j s_{ij}$$ -Subtracting $m_i$ before exponentiation makes the largest exponent input in the row zero and avoids excessively large values. The same shift applies to both the numerator and denominator, so the normalized softmax result is unchanged. The unnormalized attention weight at each position is: +Basic stable softmax uses $m_i^{\max}$ as its exponent reference. Subtracting it before exponentiation makes the largest exponent input in the row zero and avoids excessively large values. The same shift applies to both the numerator and denominator, so the normalized softmax result is unchanged. The unnormalized attention weight at each position is: -$$p_{ij}=\exp\left(\frac{s_{ij}-m_i}{\sqrt d}\right)$$ +$$p_{ij}=\exp\left(\frac{s_{ij}-m_i^{\max}}{\sqrt d}\right)$$ Summing the $p_{ij}$ values in the row gives the unnormalized weight sum $\ell_i$. Using the same $p_{ij}$ values to weight the value vectors gives an output vector $o_i$ that has not yet been divided by $\ell_i$: @@ -51,7 +51,7 @@ The final output is: $$O_i=\frac{o_i}{\ell_i}$$ -FlashAttention processes K/V in blocks. Once a block's scores have been consumed, they can be discarded; later blocks need only the running $m_i$, $\ell_i$, and $o_i$. Both $\ell_i$ and $o_i$ were accumulated using the reference $m_i$ in effect at the time. If a later block adopts a larger reference, the old state must first be converted to the new scale before the current block's contribution can be added. +FlashAttention processes K/V in blocks. Once a block's scores have been consumed, they can be discarded. For each row, the kernel retains an exponent reference $r_i$, the running denominator $\ell_i$, and the running weighted sum $o_i$. Basic online softmax updates $r_i$ to the largest score seen so far, whereas FA4 may temporarily keep an older value. Both $\ell_i$ and $o_i$ are accumulated relative to the current $r_i$. If a later block adopts a larger reference, the old state must first be converted to the new scale before the current block's contribution can be added. Basic online softmax performs this conversion whenever it encounters a larger row maximum. FA4 first checks the gap between the old and candidate references. When the gap is small enough, it retains the old reference and avoids immediately rescaling the accumulated output. To understand this optimization, we first derive the scale conversion caused by changing the reference. @@ -63,35 +63,35 @@ The natural exponential can then be written as: $$\exp\left(\frac{s-m}{\sqrt d}\right)=2^{(s-m)\alpha}$$ -The code calls $\alpha$ `scale_log2`. Let $m_{\mathrm{old}}$ be the reference used by the running state and $m_{\mathrm{block}}$ be the row maximum of the current block. With $c$ denoting the candidate, the candidate reference is: +The code calls $\alpha$ `scale_log2`. Let $r_{\mathrm{old}}$ be the reference used by the running state and $m_{\mathrm{block}}$ be the row maximum of the current block. With $c$ denoting the candidate, the candidate reference is: -$$m_c=\max(m_{\mathrm{old}},m_{\mathrm{block}})$$ +$$r_c=\max(r_{\mathrm{old}},m_{\mathrm{block}})$$ Define their signed gap in the base-2 exponent domain as $\delta$, corresponding to the code variable `delta`: -$$\delta=(m_{\mathrm{old}}-m_c)\alpha\le 0$$ +$$\delta=(r_{\mathrm{old}}-r_c)\alpha\le 0$$ -$\delta$ is the old reference minus the candidate reference, measured in base-2 exponent units. Thus $-\delta$ is the amount by which the candidate exceeds the old reference. Because $m_c\ge m_{\mathrm{old}}$, $\delta$ cannot be positive. +$\delta$ is the old reference minus the candidate reference, measured in base-2 exponent units. Thus $-\delta$ is the amount by which the candidate exceeds the old reference. Because $r_c\ge r_{\mathrm{old}}$, $\delta$ cannot be positive. The [FA4 paper](https://arxiv.org/abs/2603.05451) typically sets the threshold to $\tau=\log_2(256)=8$. When $-\delta=8$, retaining the old reference lets the largest unnormalized weight in the current block reach $2^8=256$; switching to the candidate reference would instead multiply the old state by $2^\delta=1/256$. The threshold therefore permits at most a 256-fold scale gap before rescaling: `delta >= -8` retains the old reference, whereas `delta < -8` changes the reference. Using this threshold to delay rescaling reduces the data movement and multiplications performed by the correction path; the value 8 balances fewer rescaling operations against bounded exponent growth. -If this iteration adopts the candidate reference $m_c$, every exponential accumulated under the old reference must be multiplied by the same factor: +If this iteration adopts the candidate reference $r_c$, every exponential accumulated under the old reference must be multiplied by the same factor: -$$e^{(s-m_c)/\sqrt d} -=e^{(s-m_{\mathrm{old}})/\sqrt d} -\cdot e^{(m_{\mathrm{old}}-m_c)/\sqrt d}$$ +$$e^{(s-r_c)/\sqrt d} +=e^{(s-r_{\mathrm{old}})/\sqrt d} +\cdot e^{(r_{\mathrm{old}}-r_c)/\sqrt d}$$ Writing this conversion factor as $a_{\mathrm{scale}}$ gives: $$a_{\mathrm{scale}} -=e^{(m_{\mathrm{old}}-m_c)/\sqrt d} +=e^{(r_{\mathrm{old}}-r_c)/\sqrt d} =2^\delta$$ -After switching to the candidate reference $m_c$, the accumulated denominator $\ell_i$ and weighted sum $o_i$ remain on the old scale. The kernel first multiplies both by $a_{\mathrm{scale}}=2^\delta$ to convert them to the new scale, then adds the current block's contributions. In the pseudocode below, $\ell_i$ and $o_i$ become `row_sum` and `O`, while `acc_scale = exp2(delta)` computes the conversion factor. +After switching to the candidate reference $r_c$, the accumulated denominator $\ell_i$ and weighted sum $o_i$ remain on the old scale. The kernel first multiplies both by $a_{\mathrm{scale}}=2^\delta$ to convert them to the new scale, then adds the current block's contributions. In the pseudocode below, $\ell_i$ and $o_i$ become `row_sum` and `O`, while `acc_scale = exp2(delta)` computes the conversion factor. The three values retained across K/V blocks map to the pseudocode as follows: -- `row_max`: the reference subtracted from every score in the row before exponentiation, namely $m_i$. The basic algorithm uses the largest score seen so far; FA4 may retain the old value while the threshold permits it. +- `row_max`: the exponent reference $r_i$ subtracted from every score in the row. Basic online softmax uses the largest score seen so far; FA4 may retain the old reference while the threshold permits it. Despite its name, `row_max` therefore need not equal the exact maximum $m_i^{\max}$ at every iteration. - `row_sum`: the sum of $p_{ij}$ over all key positions processed so far, namely $\ell_i$. - `O`: the weighted sum $o_i$ formed from the same $p_{ij}$ values. It is divided by `row_sum` only after all blocks have been processed. diff --git a/chapter_gemm_advanced/index.md b/chapter_gemm_advanced/index.md index a60bcd5f..172af2b6 100644 --- a/chapter_gemm_advanced/index.md +++ b/chapter_gemm_advanced/index.md @@ -22,7 +22,7 @@ This chapter broadens cooperation in three steps. Step 7 assigns TMA, MMA, and w In the single-warpgroup kernel, every thread follows the same load, compute, and writeback path. The Tensor Cores have no work while data is being loaded, and the TMA engine may sit idle during computation. Warp specialization assigns these jobs to different warps and uses a software pipeline to pass data between them, allowing several stages to run concurrently. > **Step 7 execution structure** -> - Scope: one warpgroup walking load → MMA → writeback in order becomes three concurrent roles (TMA producer, MMA consumer, writeback) connected by full/empty barriers. +> - Scope: the sequential load → MMA → writeback path in one warpgroup becomes three concurrent roles (TMA producer, MMA consumer, and writeback) connected by full/empty barriers. > - Layout: unchanged, same SMEM stages and TMEM accumulator as Step 6. > - Dispatch: unchanged, TMA loads, `tcgen05` MMA. @@ -57,7 +57,7 @@ With `WG_NUMBER=2`, the kernel uses two warpgroups and assigns load, compute, an ### Four Barriers -The three concurrent roles communicate through four barriers. The forward path, TMA → MMA → Writeback, reports that data is ready. The reverse path, Writeback → MMA → TMA, releases a buffer for reuse. Barrier names follow `source2destination`; for example, `tma2mma` carries the notification from TMA to MMA. +The three concurrent roles communicate through four barriers. The forward path, TMA → MMA → Writeback, reports that data is ready. The reverse path, Writeback → MMA → TMA, returns each protected buffer or resource to the preceding role for reuse. Barrier names follow `source2destination`; for example, `tma2mma` carries the notification from TMA to MMA. | Barrier | Type | Direction | Meaning | |---------|------|-----------|---------| @@ -876,7 +876,7 @@ The table below follows the progression from the naive baseline to the warp-spec Every row with a measured time uses the same `M=N=K=4096` problem, so those rows can be compared directly. The 70 ms in Step 1 comes from a full-matrix baseline with the same sequential data path; it is not a run of the single-tile `hgemm_v1` from {ref}`chap_gemm_basics`. The introductory chapter uses smaller problems to explain Steps 1 through 3, while the Step 1 and Step 3 rows here measure the corresponding full-matrix implementations. -Step 2 still computes only one output tile, so it is not directly comparable with the full-matrix results. Steps 5 and 6 are intermediate versions between the TMA-load kernel and the warp-specialized kernel; their mechanisms are retained in Step 7. The table therefore shows only the endpoints of that interval. Steps 2, 5, and 6 use dashes and have no standalone speedup. +Step 2 still computes only one output tile, so it is not directly comparable with the full-matrix results. Steps 5 and 6 are intermediate versions between the TMA-load kernel and the warp-specialized kernel; their mechanisms are retained in Step 7. The table therefore shows only the endpoints of that interval. Steps 2, 5, and 6 use dashes, so no cumulative speedup relative to Step 1 is shown for them. These numbers come from one B200 reference run. They are intended to compare the versions in this tutorial under the same conditions, rather than to represent peak performance for other problem sizes or environments. diff --git a/chapter_gemm_basics/index.md b/chapter_gemm_basics/index.md index f9878dfa..372217c2 100644 --- a/chapter_gemm_basics/index.md +++ b/chapter_gemm_basics/index.md @@ -628,6 +628,6 @@ def hgemm_v3(M, N, K): ## Exercises -1. In Steps 1-3, `Tx.copy` moves A and B tiles into SMEM before MMA. Why does the kernel need `T.cuda.cta_sync()` before `Tx.gemm_async` reads those SMEM tiles? +1. In Steps 1-3, `Tx.cta.copy` moves A and B tiles into SMEM before MMA. Why does the kernel need `T.cuda.cta_sync()` before `Tx.gemm_async` reads those SMEM tiles? 2. In Step 2, what happens if `phase_mma ^= 1` is removed from the K-loop? Does the kernel wait for every MMA, or can a later wait pass too early? 3. For `M=N=4096` and `BLK_M=BLK_N=128`, what is the Step 3 grid shape, and how many CTAs does it launch? For CTA `(bx, by)`, which other CTAs independently read the same A tiles, and which independently read the same B tiles? Does the current kernel explicitly share that data? diff --git a/chapter_layout_generations/index.md b/chapter_layout_generations/index.md index e1e766ce..7cbd60e9 100644 --- a/chapter_layout_generations/index.md +++ b/chapter_layout_generations/index.md @@ -266,7 +266,7 @@ S[(4, 32, 4) : (4@TCol, 1@TLane, 1@TCol)] `S[...]` maps `(Mgroup, lane, sfk)` to byte positions in TMEM. In a typed TIRx layout, `@TCol` strides are measured in buffer elements. A scale factor is 8 bits here, so the logical TCol position is `4*Mgroup+sfk`; every four consecutive positions pack into one 32-bit hardware TCol cell. Equivalently, `hardware_TCol=(4*Mgroup+sfk)//4` and `byte_in_word=(4*Mgroup+sfk)%4`. -`R[...]` adds four replicas along `TLane`. The `.32x128b.warpx4` form of `tcgen05.cp` creates this layout: it writes one 32-lane window, then broadcasts the same data into the other three warp windows. +`R[...]` adds four replicas along `TLane`. The `.32x128b.warpx4` form of `tcgen05.cp` creates this layout by multicasting the same base tile into four 32-lane warp windows. ### Word-Level Replication for `scale_vec` diff --git a/chapter_performance/index.md b/chapter_performance/index.md index 37aad646..f230b5be 100644 --- a/chapter_performance/index.md +++ b/chapter_performance/index.md @@ -139,7 +139,7 @@ Once a kernel is known to be memory-bound, there are two avenues for optimizatio traffic to raise arithmetic intensity, or, when the traffic cannot be reduced further, bring effective bandwidth as close as possible to the hardware limit. -Fusion is often the most direct method. A common source of low arithmetic intensity is that one kernel writes an intermediate tensor to HBM, and the next operation immediately reads it back. After fusing the producer, which creates the intermediate, with the consumer, which uses it, the intermediate can stay in registers or on-chip storage such as SMEM or TMEM, avoiding that HBM round trip. +Fusion is often the most direct method. A common source of low arithmetic intensity is an intermediate tensor that one kernel writes to HBM and the next operation immediately reads back. Fusing the operation that produces the intermediate with the operation that consumes it can keep the value in registers or on-chip storage such as SMEM or TMEM, avoiding the HBM round trip. - Fuse GEMM with an elementwise epilogue. - Fuse normalization into an adjacent operator. @@ -238,7 +238,7 @@ implementation that approaches that ceiling. A large fp16 GEMM may be compute-bound in theory. That only means the HBM-level memory roof is not the main limit; it does not mean any implementation will reach the Tensor Core compute roof. Closing the gap requires the right instructions, layouts, staging, synchronization, and scheduling. The later GEMM chapters show this on B200 through a sequence of steps: each step keeps the same basic algorithm but changes how the tile is computed or scheduled. -In the GEMM optimization ladder, the first large measured jump is the move from the thread-copy tiled path to the TMA-backed path. The former uses ordinary CTA threads to copy tiles from GMEM to SMEM; the latter delegates this regular tile movement to the TMA hardware engine, letting the kernel feed Tensor Cores through hardware-managed bulk copies. +In the GEMM optimization ladder, the first large measured jump is the move from the thread-copy tiled path to the TMA-backed path. The former uses ordinary CTA threads to copy tiles from GMEM to SMEM; the latter delegates this regular tile movement to the TMA hardware engine. TMA fills SMEM through hardware-managed bulk copies, and the MMA path then reads those tiles from SMEM. After that first jump, subsequent optimizations address one question: how can the kernel reduce waiting among data movement, Tensor Core computation, and the epilogue? Software pipelining and warp @@ -257,7 +257,7 @@ CTA clusters, and multi-consumer execution each change. ## Reducing Idle Time Through Overlap -Once a GEMM is compute-bound and already uses Tensor Cores, the remaining gap usually comes from hardware idle time. +Once a GEMM is compute-bound and already uses Tensor Cores, the remaining gap usually reflects periods when one or more execution paths are not fully utilized. A simple kernel might do this: @@ -282,7 +282,7 @@ store tile k - 1 ``` On Blackwell, TMA, `tcgen05.mma`, and the epilogue/store path primarily execute these three stages, -while `mbarrier` coordinates data handoffs among them. +while `mbarrier` coordinates completion and buffer ownership between them. Overlap does not remove dependencies. The MMA for tile `k` must still wait for that tile to load, and the epilogue must still wait for the MMA to complete. The kernel can instead advance independent @@ -302,7 +302,7 @@ each CTA uses a large amount of shared memory, fewer CTAs or warps fit on the SM Many modern Tensor Core kernels intentionally spend resources in ways that reduce occupancy. Multi-stage shared memory pipelines consume SMEM. Large register fragments consume registers. TMEM allocations consume Tensor Memory capacity. Warp specialization may reserve whole warps for producer or consumer roles. This is a deliberate tradeoff. Rather than hiding latency with many resident warps, these kernels -explicitly overlap stages within a smaller number of resident CTAs. A low-occupancy kernel can still +explicitly overlap stages within a smaller number of resident CTAs. A low-occupancy kernel may still perform well if its pipeline keeps TMA, Tensor Cores, and the store path active. Each approach has its place. Kernels with irregular memory access or limited opportunities for an diff --git a/chapter_tensor_cores/index.md b/chapter_tensor_cores/index.md index 67829e8a..44dcf866 100644 --- a/chapter_tensor_cores/index.md +++ b/chapter_tensor_cores/index.md @@ -113,7 +113,7 @@ With `cta_group::1`, the MMA updates only the current CTA's TMEM. We begin with With `cta_group::2`, the MMA accesses the TMEM of both CTAs in a pair. A CTA pair consists of two CTAs in the same cluster whose `%cluster_ctarank` values differ only in the least-significant bit. One rank is even and the other is odd; we refer to them below as the even CTA and the odd CTA. -Only one thread in the CTA pair needs to issue `tcgen05.mma`. That thread may belong to either CTA, but the peer CTA must remain active. The kernels later in this book generally elect one thread in the even CTA to issue the MMA and use `tcgen05.commit` to arrange completion notification. +Only one thread in the CTA pair needs to issue `tcgen05.mma`. That thread may belong to either CTA, but the peer CTA must remain active. The kernels later in this book generally elect one thread in the even CTA to issue the MMA and use `tcgen05.commit` to associate the completion of the previously issued operations with an `mbarrier`. The accumulator layout depends on four choices: `cta_group`, the size of M, whether A is dense or structured sparse, and whether the instruction is ordinary `tcgen05.mma` or the weight-stationary `tcgen05.mma.ws`. The selected layout maps each logical coordinate `(m,n)` to `TLane` and `TCol`. diff --git a/chapter_tirx_layout_api/index.md b/chapter_tirx_layout_api/index.md index 5cb60ff7..e6ae60c3 100644 --- a/chapter_tirx_layout_api/index.md +++ b/chapter_tirx_layout_api/index.md @@ -142,7 +142,7 @@ R[2 : 4@warpid] places two copies along the `warpid` axis, separated by four warps. -GPU hardware often needs to broadcast the same data across warps, lanes, or storage regions. A replica expresses that behavior directly as one logical element with several physical coordinates. +A replica describes one logical element as having several physical coordinates. It records where the copies belong; the tile operation that consumes the layout determines how those copies are produced or used. ### Offset @@ -406,7 +406,7 @@ Logical element `(i, j)` has linear element address: m = 64 * i + j ``` -Each row contains 64 float16 values, or 128 bytes. Reading a fixed column `j` advances by 128 bytes from one row to the next, so several accesses may repeatedly land in the same set of banks. +Each row contains 64 float16 values, or 128 bytes. If a group of threads reads the same column `j` from different rows, successive addresses are 128 bytes apart and may repeatedly land in the same set of banks. A swizzle makes low address bits depend on higher row bits, scattering a column access that would otherwise repeatedly hit the same bank. diff --git a/index.md b/index.md index d8da68ba..70b264a8 100644 --- a/index.md +++ b/index.md @@ -1,26 +1,28 @@ # Modern GPU Programming For MLSys -Machine learning systems sit at the heart of modern AI workloads. In these systems, performance -often comes down to the quality of a small number of GPU kernels. Attention kernels, LLM prefill -and decode kernels, low-precision block-scaled GEMMs, fused MoE layers, and other large fused -kernels all directly shape end-to-end speed in both training and serving. - -To make these kernels fast, however, we need more than a list of optimization tricks. Modern GPUs -are no longer simple variations of the same old design. Recent architectures introduce richer -memory spaces, new access patterns, and increasingly specialized execution units. To program them -well, we need both a clear mental model of the hardware and a practical understanding of how -high-performance kernels are built. This book is about developing both. - -The book follows a simple progression: first understand the GPU hardware, then learn the -programming model we will use, and finally build state-of-the-art kernels step by step. Our main -target is the Blackwell generation, and our main running examples are General Matrix-Matrix -Multiplication (GEMM) and FlashAttention. Along the way, we will also study the core ingredients -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. To make the ideas easier to study and easier to run, this book uses -the **TIRx** Python DSL to build real GPU kernel examples step by step. TIRx stays close to the -hardware, which lets us reason about low-level control while still learning through runnable code. +Machine learning systems power many of today's AI workloads. As models grow and deployment +settings become more complex, end-to-end performance increasingly depends on a small number of +critical GPU kernels. Attention, LLM prefill and decode, low-precision block-scaled GEMM, fused MoE +layers, and other large fused kernels directly affect both training and serving speed. + +Making these kernels fast requires more than a list of optimization tricks. Recent GPU +architectures introduce richer memory spaces, new data-movement mechanisms, and increasingly +specialized execution units. Using them effectively requires both a clear understanding of how the +hardware executes a program and practical knowledge of how a basic kernel evolves into a +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. + +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 +studied, run, and verified in real kernels. TIRx keeps hardware-level choices explicit, making it +possible to reason about control flow, memory access, and synchronization while working with +runnable code. This book is open source. Contributions, corrections, and examples are welcome through the [GitHub repository](https://github.com/mlc-ai/modern-gpu-programming-for-mlsys). @@ -28,7 +30,7 @@ This book is open source. Contributions, corrections, and examples are welcome t ## How This Book Is Organized - **Part I, Understanding the GPU.** This part introduces the overall organization of the GPU, - general recipes for writing fast kernels, and key concepts such as data layout, asynchronous + general techniques for writing fast kernels, and key concepts such as data layout, asynchronous memory operations, and coordination. It builds the hardware intuition that the rest of the book relies on. - **Part II, TIRx Overview.** This part introduces the key elements of TIRx, which serve as the @@ -37,7 +39,8 @@ This book is open source. Contributions, corrections, and examples are welcome t 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. -- **Reference.** TIRx language reference and compiler internals. +- **Reference.** TIRx language reference, compiler internals, and a guide to debugging asynchronous + kernels. ```{toctree} :caption: Part I, Understanding the GPU diff --git a/tirx_guide/arch/lowering_pipeline.rst b/tirx_guide/arch/lowering_pipeline.rst index 3b19731b..852ee618 100644 --- a/tirx_guide/arch/lowering_pipeline.rst +++ b/tirx_guide/arch/lowering_pipeline.rst @@ -18,11 +18,11 @@ TIRx lowering pipeline ====================== -``tvm.compile(mod, target, tir_pipeline="tirx")`` runs an authored TIRx module -through the **tirx pipeline** — an ordered sequence of TIR passes that turns the -high-level constructs you write (tile primitives, ``TileLayout``-typed buffers, -execution-scope ids) into split **host** + **device** functions, which the CUDA -backend then renders to source. The pipeline is defined in +``tvm.compile(mod, target, tir_pipeline="tirx")`` passes a TIRx module through the +**tirx pipeline**, an ordered sequence of TIR passes. These passes lower the +high-level constructs in the source—tile primitives, ``TileLayout``-typed buffers, +and execution-scope IDs—and split the module into **host** and **device** functions. +The CUDA backend then generates source for the device functions. The pipeline is defined in ``python/tvm/tirx/compilation_pipeline.py`` (``tirx_pipeline``); this page walks the passes in order. @@ -128,9 +128,9 @@ Inside LowerTIRx LowerTIRx = Sequential([ TilePrimitiveDispatch, LowerTIRxCleanup ]) -- **``TilePrimitiveDispatch``** replaces every ``TilePrimitiveCall`` (``copy``, - ``gemm``, ``reduction``, …) with the body emitted by its selected backend - dispatch — its variant-selection and codegen. +- **``TilePrimitiveDispatch``** selects a backend variant for every + ``TilePrimitiveCall`` (``copy``, ``gemm``, ``reduction``, …) and replaces the + call with the implementation emitted by that variant. - **``LowerTIRxCleanup``** runs the ``LayoutApplier``: it resolves every ``TileLayout``-typed buffer access into concrete physical address arithmetic (``addr = data + elem_offset + layout.apply(coord)``), flattens the buffers, and diff --git a/tirx_guide/language_reference/cuda/buffers.rst b/tirx_guide/language_reference/cuda/buffers.rst index 2ef2820a..87983caf 100644 --- a/tirx_guide/language_reference/cuda/buffers.rst +++ b/tirx_guide/language_reference/cuda/buffers.rst @@ -240,9 +240,10 @@ The TMEM pool (`Tensor memory`_, below) is layered on top of an ``SMEMPool``. Registers --------- -Per-thread scratch lives in registers. Allocate it with ``T.alloc_local(shape, -dtype)`` (i.e. ``scope="local"``): it is private to each thread and lowers to a -local array kept in registers. +Per-thread scratch uses ``local`` scope. Allocate it with ``T.alloc_local(shape, +dtype)`` (i.e. ``scope="local"``): it is private to each thread. Statically indexed +local arrays are normally scalarized into registers, while dynamically indexed +arrays or values under high register pressure may use local memory. .. code-block:: python @@ -376,7 +377,7 @@ tensor as a view at a column offset, and one warp frees it at the end: # ... use acc as a gemm_async / copy_async operand ... if warp_id == alloc_warp: T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) - T.ptx.tcgen05.dealloc(addr, n_cols=512, cta_group=cta_group) + T.ptx.tcgen05.dealloc(addr[0], n_cols=512, cta_group=cta_group) You manage the column offsets and the ``tmem_layout`` (a datapath D/F layout) yourself. This is exactly the sequence the pool below emits. diff --git a/tirx_guide/language_reference/cuda/control_flow.rst b/tirx_guide/language_reference/cuda/control_flow.rst index e30849de..59d304bf 100644 --- a/tirx_guide/language_reference/cuda/control_flow.rst +++ b/tirx_guide/language_reference/cuda/control_flow.rst @@ -18,8 +18,8 @@ Control flow ============ -Control flow is ``if``, the loop family, and ``while`` — each maps to the obvious -CUDA. +TIRx provides ``if``, several loop forms, and ``while``. Each maps directly to the +corresponding CUDA control flow. if -- @@ -46,8 +46,9 @@ thread/lane comparison, or elect a single issuing thread with A_ptr[tx] = A_ptr[tx] + 1.0f; } -For an expression-level choice (no branch), use ``T.if_then_else(cond, a, b)``. It -lowers to a ternary, so it introduces no control-flow divergence: +For an expression-level choice without an explicit TIRx control-flow branch, use +``T.if_then_else(cond, a, b)``. It lowers to a ternary expression; the backend +still decides which machine instructions implement that expression: .. code-block:: c++ @@ -67,9 +68,11 @@ warpgroups will never arrive and the kernel will deadlock. When only one warpgro needs to synchronize, use a warpgroup-scoped ``T.cuda.warpgroup_sync(id)`` (see :ref:`chap_gemm_advanced` and :doc:`threads_sync`). -The same caution applies to barrier setup. An ``mbarrier`` ``.init()`` lowers to a -single-thread guard (``if (threadIdx.x < 1)``). Nesting it inside another divergent -branch can leave the barrier uninitialized, leading to unspecified launch failures. +The same caution applies to barrier setup. The high-level ``MBarrier.init()`` +wrapper emits a single-thread guard (``if (threadIdx.x < 1)``). Nesting that call +inside another divergent branch can leave the barrier uninitialized and cause an +unspecified launch failure. The raw ``T.ptx.mbarrier.init`` intrinsic does not add +this guard automatically; its caller must select the initializing thread. loop ---- diff --git a/tirx_guide/language_reference/cuda/data_types.rst b/tirx_guide/language_reference/cuda/data_types.rst index 472ac8f0..3d2f8565 100644 --- a/tirx_guide/language_reference/cuda/data_types.rst +++ b/tirx_guide/language_reference/cuda/data_types.rst @@ -66,7 +66,7 @@ declares a ``float4`` register directly (you index it as ``v[0]``), and a ``float32x4`` ``vload`` / ``vstore`` then moves it as one 16-byte access. The vector dtype is not tied to ``vload`` — any buffer or scalar can carry it. -so the dtype → CUDA mapping is: +The resulting dtype → CUDA mapping is: .. list-table:: :header-rows: 1 diff --git a/tirx_guide/language_reference/cuda/threads_sync.rst b/tirx_guide/language_reference/cuda/threads_sync.rst index c7783147..7ce083b0 100644 --- a/tirx_guide/language_reference/cuda/threads_sync.rst +++ b/tirx_guide/language_reference/cuda/threads_sync.rst @@ -18,8 +18,8 @@ CUDA C++/PTX intrinsics ======================= -When no tile primitive covers what you need, two escape hatches reach the hardware -directly: **call a backend intrinsic** (the ``T.cuda.*`` / ``T.ptx.*`` namespaces +When no tile primitive covers the required operation, two lower-level options are +available: **call a backend intrinsic** (the ``T.cuda.*`` / ``T.ptx.*`` namespaces from ``tvm.backend.cuda``), or **inline raw CUDA** source. Calling backend intrinsics @@ -78,9 +78,10 @@ misusing any of them usually leads to silent corruption or deadlock. The ``T.ptx.mbarrier.try_wait(bar, phase)`` intrinsic blocks until the barrier's internal phase *differs* from the ``phase`` argument provided by the caller. Consequently, when reusing a barrier across loop iterations, the caller must flip -its local phase tracker (``phase ^= 1``) after every wait. Failing to do so causes -subsequent waits to return immediately, allowing the engine to read half-written -memory. :ref:`chap_gemm_basics` walks through the full phase-tracking table. +its local phase tracker (``phase ^= 1``) after every wait. Failing to do so can +cause a later wait to return for an earlier phase, allowing a consumer to access +data before the current producer or asynchronous operation has finished. +:ref:`chap_gemm_basics` walks through the full phase-tracking table. **Election.** ``T.ptx.elect_sync()`` elects a *single active lane within a warp*, not lane 0, and not one thread per CTA. To narrow an issuer down to exactly one @@ -93,9 +94,10 @@ requires *every* CTA thread to arrive. Once warpgroups specialize onto different code paths, placing a ``cta_sync()`` inside a warpgroup branch deadlocks the kernel because the other warpgroups never reach it. The hardware provides 16 named barriers (IDs 0 to 15); ``T.cuda.warpgroup_sync(10)`` synchronizes only the threads -of one warpgroup. Distinct warpgroups take distinct IDs (e.g., -``warpgroup_sync(wg_id + 10)``) so they never collide on the same hardware barrier. -See :ref:`chap_gemm_advanced`. +of one warpgroup. Independent synchronizations that may be active at the same time +must use distinct IDs (for example, ``warpgroup_sync(wg_id + 10)``). An ID may be +reused after the preceding synchronization that used it has completed. See +:ref:`chap_gemm_advanced`. **Fences.** Fences order a producer's writes before a consumer (often an asynchronous engine) reads them: diff --git a/zh/appendix/debugging_warp_specialized.md b/zh/appendix/debugging_warp_specialized.md index 47cacd93..d7dfc33b 100644 --- a/zh/appendix/debugging_warp_specialized.md +++ b/zh/appendix/debugging_warp_specialized.md @@ -1,12 +1,13 @@ (chap_warp_spec_debug)= -# 调试 Warp-Specialized Kernel +# 调试 Warp-Specialized Kernels -第三部分 GEMM 的第 7-9 步让 TMA load、`tcgen05` MMA 和 TMEM/SMEM writeback -重叠执行。Flash Attention 也使用类似的数据交接方式。调试这类 kernel 时， -可以先确定各个角色及其使用的存储空间，再检查生成的 CUDA 是否符合这个模型。 +第三部分 GEMM 的第 7 至第 9 步让 TMA load、`tcgen05` MMA 和 TMEM/SMEM +writeback 重叠执行。Flash Attention 中的 QKᵀ MMA、softmax、PV MMA 和 +correction 也使用类似的交接方式。调试这类 kernel 时，可以先确定各个角色及其 +使用的存储空间，再检查生成的 CUDA 是否符合这个模型。 -不要一开始就重写 kernel。先确认运行环境正确，并用最小测试稳定复现问题， -然后再检查生成的 CUDA。排除环境和编译问题后，这类 kernel 的运行时故障通常 +不要一开始就重写 kernel。先确认运行环境无误，并用最小的正确性测试稳定复现 +问题，再检查生成的 CUDA。排除环境和编译问题后，这类 kernel 的运行时故障通常 来自某次数据交接：barrier 没有初始化、arrival count 错误、collective 的参与 范围被角色分支缩小、wait 使用了旧的 barrier phase，或者 producer 的写入尚未 可见，存储空间就被提前复用。 @@ -22,11 +23,11 @@ python -c "import torch; print(torch.cuda.get_device_name(), torch.cuda.get_devi 这些 kernel 面向 Blackwell（`sm_100a`）。如果 Python 导入了旧的 TVM checkout，或者当前 GPU 不是 Blackwell 架构，应先修正环境，再修改 kernel。 -随后先运行最小的正确性测试，例如 `run_correctness()`；正确性通过后再看性能。 +环境确认无误后，先运行最小的正确性测试，例如 `run_correctness()`；正确性通过后再看性能。 ## 调试步骤 -1. 用仍能复现问题的最小 shape 运行。如果发生 illegal memory access，下一次运行前先重启 Python。 +1. 将输入缩小到仍能稳定复现问题的最小 shape。如果发生 illegal memory access，下一次运行前先重启 Python。 2. 如果编译失败，先检查已安装的 API、target、`dispatch=` 和 buffer scope，再检查运行时同步代码。 3. 保存 `inspect_source("cuda")` 的输出。先搜索 role guard、`mbarrier_init`、`tcgen05`、`cp.async.bulk.tensor` 和 `cta_sync()`，再回头阅读 Python。 4. 针对出错的 kernel 路径，写出 roles、storage、handoff 和 lifetime 表。 @@ -52,20 +53,20 @@ checkout，或者当前 GPU 不是 Blackwell 架构，应先修正环境，再 - Barrier 初始化出现在各个角色分支之前。 - Collective 没有被 lane、warp 或 warpgroup guard 意外缩小参与范围。 - Arrive/wait phase 与 handoff 表一致。 -- 完成相应的 wait，并且 lifetime 表允许之后，才释放 TMEM 或复用 SMEM。 +- 必须确认 TMA store 已经完成，并且 lifetime 表表明相关资源可以复用，之后才能释放 TMEM 或复用相应的 SMEM。 -这张表既适用于 GEMM 的 TMA -> MMA -> writeback pipeline，也适用于 -Flash Attention 中 score、softmax、value 和 correction 之间的交接。 +这张表既适用于 GEMM 的 TMA → MMA → writeback pipeline，也适用于 +Flash Attention 中 QKᵀ MMA、softmax、PV MMA 和 correction 之间的交接。 ## 编译失败 先解决编译问题，再调试运行时同步： -| 现象 | 可能的位置 | 首先检查 | +| 现象 | 可能原因 | 首先检查 | |---|---|---| -| TIRx API 未知或发生 attribute error | 安装的 wheel 与教程代码不匹配 | 输出 `tvm.__file__` 和 `tvm.__version__`，并对照 {ref}`chap_language_reference` 检查 API 名称。 | +| 找不到 TIRx API，或出现 attribute error | 安装的 wheel 与教程代码不匹配 | 输出 `tvm.__file__` 和 `tvm.__version__`，并对照 {ref}`chap_language_reference` 检查 API 名称。 | | 不支持指定的 `dispatch=` | 当前 target 或 primitive 不支持这条路径 | 检查 `dispatch` 参数和 target capability；本教程中的 `tcgen05` 路径需要 Blackwell。 | -| Buffer scope 不匹配 | Buffer 通过错误的硬件路径使用 | 检查表中的 storage：TMEM 必须通过 `tcgen05` 访问，TMA 搬运的 buffer 必须使用兼容的 GMEM/SMEM layout。 | +| Buffer scope 不匹配 | Buffer 被交给了不匹配的硬件路径 | 检查表中的 storage：TMEM 必须通过 `tcgen05` 访问，TMA 搬运的 buffer 必须使用兼容的 GMEM/SMEM layout。 | | 编译成功，但生成的 CUDA 中没有预期路径 | Dispatch 没有生成预期的硬件指令 | 修改算法前，先在生成的 CUDA 中搜索 `tcgen05` 和 `cp.async.bulk.tensor`。 | ## 检查生成的代码 @@ -109,25 +110,25 @@ print(cuda_source) guard；在生成的 CUDA 中，应搜索上表对应的表达式。 ```c -// (1) Barrier inits: top level, CTA thread 0 only +// (1) Barrier 初始化：位于顶层，只由 CTA thread 0 执行 if (threadIdx.x < 1) { mbarrier_init(tma2mma[0..1], 1); mbarrier_init(mma2tma[0..1], 1); mbarrier_init(mma2ld, 1); - mbarrier_init(ld2mma, 128); // arrived by all 128 WG0 threads + mbarrier_init(ld2mma, 128); // WG0 的 128 个 threads 全部执行 arrival } -// (2) TMEM alloc: WG0 warp 0, all lanes of the issuing warp +// (2) TMEM 分配：WG0 warp 0，发出指令的 warp 中所有 lanes 都参与 if (wg_id == 0 && warp_id == 0) tcgen05_alloc(..., 512); -// (3) Fences + cta_sync, then phase init: producer=1, consumer=0 +// (3) 执行 fences 和 cta_sync，再初始化 phase：producer=1，consumer=0 // (4) Warp-specialized loop if (wg_id == 1 && warp_id == 3 && elect_sync) { /* TMA */ while(valid){ ... next_tile(); } } if (wg_id == 1 && warp_id == 0 && elect_sync) { /* MMA */ while(valid){ ... next_tile(); } } if (wg_id == 0) { /* WB */ while(valid){ ... next_tile(); } } -// (5) Cleanup: issuing warp, no lane guard +// (5) 清理：由发出指令的 warp 执行，不使用 lane guard cta_sync(); if (warp_id == 0) { tcgen05_relinquish_alloc_permit(); tcgen05_dealloc(..., 512); } ``` @@ -143,7 +144,7 @@ if (warp_id == 0) { tcgen05_relinquish_alloc_permit(); tcgen05_dealloc(..., 512) 现象只能作为线索，不应直接当作最终诊断： -| 线索 | 可能的位置 | 首先检查 | +| 线索 | 可能原因 | 首先检查 | |---|---|---| | Kernel 卡住，随后 runtime 报告 unspecified launch failure | Deadlock | Barrier 初始化的位置、arrival count、`cta_sync()` 的位置和 `next_tile()` 的参与范围 | | Illegal memory access、XID，或之后无关的 CUDA 调用也失败 | Crash / poisoned context | 重启 Python，再检查 pointer 范围、storage lifetime 和 collective 的参与范围 | @@ -201,20 +202,20 @@ descriptor、operand 设置或未初始化的 accumulation。数值有限但错 - **TMA store 前缺少 `fence.proxy_async("shared::cta")`。** TMA engine 可能看不到 threads 对 SMEM 的写入。 - **TMA store 后缺少 `cp_async.bulk.commit_group()` 和 `wait_group(0)`。** Store 尚未完成，下一 tile 就复用了 Dsmem。 - **Persistent kernel 在 `1024×1024` 等较小 shape 上偶发失败。** 更大的 shape 和更长的 K-loop 可能掩盖竞争。重新检查 tiles 之间的 phase reset 和 TMA store commit/wait。 -- **等待 MMA 完成后直接读取 TMEM。** `mma2ld.wait` 只能确认 MMA 已经完成；writeback thread 在随后发出 `tcgen05.ld` 前，还需要执行 `T.ptx.tcgen05.fence.after_thread_sync()`，把这次 TMEM load 排在跨 thread 的完成通知之后。第 7-9 步都将它放在 `mma2ld.wait` 之后。这个 fence 只负责 `tcgen05` 指令之间的顺序；等待 TMA load 和让普通 thread 的 SMEM 写入对 TMA engine 可见，分别使用各自的 mbarrier 和 proxy fence 协议。 +- **等待 MMA 完成后直接读取 TMEM。** `mma2ld.wait` 只能确认 MMA 已经完成；writeback thread 在随后发出 `tcgen05.ld` 前，还需要执行 `T.ptx.tcgen05.fence.after_thread_sync()`，把这次 TMEM load 排在跨 thread 的完成通知之后。第 7 至第 9 步都将它放在 `mma2ld.wait` 之后。这个 fence 只负责 `tcgen05` 指令之间的顺序；等待 TMA load 和让普通 thread 的 SMEM 写入对 TMA engine 可见，分别使用各自的 mbarrier 和 proxy fence 协议。 ## 结果正确但性能较差 如果结果正确，但性能远低于预期，可以继续使用同一套检查流程： -| 线索 | 可能的位置 | 首先检查 | +| 线索 | 可能原因 | 首先检查 | |---|---|---| | 生成的 CUDA 中没有 `cp.async.bulk.tensor` | Copy 没有生成 TMA 路径 | 检查 `dispatch="tma"`、target capability 和 operand layout | | 生成的 CUDA 中没有 `tcgen05` | MMA 没有生成 Blackwell Tensor Core 指令 | 检查 `dispatch="tcgen05"`、target capability 和 operand layout | | TMA 与 MMA 没有重叠 | Pipeline 太浅，或者 phase 使 producer/consumer 串行执行 | 检查生成 CUDA 中 wait、arrive 和 advance 的顺序 | | 小 shape 正确，但大 shape 性能差 | Register spill、occupancy 或 staging buffer 压力 | 检查 compiler resource report；减小 tile、分块 writeback，或降低 pipeline depth | -## 提交有效的问题报告 +## 提交高质量的问题报告 如果完成上述检查后问题仍然存在，请先缩小复现范围，再到 [Apache TVM GitHub 仓库](https://github.com/apache/tvm/issues)提交 issue。 diff --git a/zh/appendix/index.md b/zh/appendix/index.md index 9690eac9..f92847cc 100644 --- a/zh/appendix/index.md +++ b/zh/appendix/index.md @@ -12,5 +12,5 @@ 完整的 `tvm.tirx` Python API 请参阅 [TVM 官方文档](https://tvm.apache.org/docs/)。 -TIRx 的基本用法见第二部分的 {ref}`chap_tirx_primer`，tensor layout 模型见 -{ref}`chap_tirx_layout_api`。 +第二部分介绍了 TIRx 编程模型（{ref}`chap_tirx_primer`）和 tensor layout 模型 +（{ref}`chap_tirx_layout_api`）。 diff --git a/zh/chapter_async_barriers/index.md b/zh/chapter_async_barriers/index.md index afc29890..f0101b49 100644 --- a/zh/chapter_async_barriers/index.md +++ b/zh/chapter_async_barriers/index.md @@ -89,7 +89,7 @@ Iteration 2 回到 stage 0，这次等待 phase 1；完成后，stage 0 的 barr `phase_tma` 描述的是软件访问 circular buffer 的轮次，与两次 TMA 传输在硬件上谁先完成无关。因此，深度为 `S` 的 TMA pipeline 通常为每个 stage 准备一个记录 TMA 完成状态的 `full` barrier，再通过 phase parity 区分同一个 stage 的前后两轮。完整的 buffer 复用协议还需要下文介绍的 `empty` barriers。 -## 常见的同步规则 +## 常见的数据交接 在 Tensor Core kernel 中，`mbarrier` 主要用于协调以下三种数据交接。 diff --git a/zh/chapter_background/index.md b/zh/chapter_background/index.md index 8691cca3..76d65c6e 100644 --- a/zh/chapter_background/index.md +++ b/zh/chapter_background/index.md @@ -65,7 +65,7 @@ TMEM 需要由程序显式管理。Kernel 必须分配和释放 TMEM；MMA 完 一个 cluster 可以包含位于不同 SM 上的多个 CTA。每个 CTA 仍然拥有自己的 shared memory，但 distributed shared memory（DSMEM）允许同一 cluster 内的其他 CTA 访问其中的数据。 -这种能力可以避免不必要的 GMEM 往返。一个 CTA 可以直接访问另一个 CTA 的 SMEM，而不需要让对方先写回 GMEM、再重新读取。使用异步操作搬运这些数据时，completion barrier 会在搬运完成后通知后续计算继续执行。 +这种能力可以避免不必要的 GMEM 往返。一个 CTA 可以直接访问另一个 CTA 的 SMEM，而不需要让对方先写回 GMEM、再重新读取。使用异步操作搬运这些数据时，操作完成后会更新 completion barrier；consumer 必须等待 barrier 完成，才能使用结果。 下图展示了一个 2-CTA cluster 中的 DSMEM 访问路径。每个 CTA 仍然拥有自己的 SMEM，但可以读取另一个 CTA 的 SMEM。 @@ -80,7 +80,7 @@ TMEM 需要由程序显式管理。Kernel 必须分配和释放 TMEM；MMA 完 在图中的 2-CTA GEMM 中，每个 CTA 都保存自己的 A 和 B 分片，同时通过 DSMEM 读取另一个 CTA 的 B 分片。这里的“共享”并不表示两个 CTA 的 SMEM 被合并成一块；它只表示 cluster 内的 CTA 可以跨 SM 访问对方的数据。 -在此基础上，两个 CTA 还可以组成 `cta_group=2`，共同执行 cooperative MMA，生成一个更大的输出 tile。 +在此基础上，两个 CTA 还可以组成 CTA pair，并以 `cta_group::2` 模式执行 cooperative MMA，生成一个更大的输出 tile。 ## 计算核心：CUDA Core 和 Tensor Core @@ -93,7 +93,7 @@ Tensor Core 的算术吞吐量远高于 CUDA Core，通常可以达到后者 10 不同 GPU 架构不仅改变 Tensor Core 的吞吐量，也改变它们的编程方式和 accumulator 的存放位置。Hopper 引入了异步 warpgroup MMA（`wgmma.mma_async`）；Blackwell 的第五代 Tensor Core，也就是 `tcgen05`，则把 accumulator 放入 Tensor Memory，而不是寄存器中。后续章节会专门讨论这一点。 -Cluster 在 GEMM 中还会带来两个重要用法。**2-CTA cooperative MMA** 允许两个 CTA 各自提供一部分 SMEM operand，共同发起一个更大的 Tensor Core MMA tile。**TMA multicast** 允许一次 GMEM load 把同一个 tile 送到多个 CTA，避免每个 CTA 分别读取同一份数据造成冗余 global memory traffic。二者都依赖前面介绍的 cluster 和 DSMEM 机制。 +Cluster 让 GEMM 可以采用两种重要的协作方式。**2-CTA cooperative MMA** 允许两个 CTA 各自提供一部分 SMEM operand，共同完成一个更大的 Tensor Core MMA tile。**TMA multicast** 允许一次 GMEM load 把同一个 tile 送到多个 CTA，避免每个 CTA 分别读取同一份数据造成冗余 global memory traffic。二者都依赖前面介绍的 cluster 和 DSMEM 机制。 ## GEMM 数据流水线 @@ -111,7 +111,7 @@ Cluster 在 GEMM 中还会带来两个重要用法。**2-CTA cooperative MMA** 单个 GEMM tile 通常会经过三个阶段。 -1. **Load：** TMA copy 把 A 或 B 的 operand tile 从 GMEM 搬到 SMEM。一个 thread 发起这次 copy，并记录预计到达的字节数。随着数据写入 SMEM，TMA 引擎会更新进度；所有预期字节到达后，completion barrier 才会被触发。 +1. **Load：** TMA copy 把 A 或 B 的 operand tile 从 GMEM 搬到 SMEM。一个 thread 发起这次 copy，并记录预计到达的字节数。随着数据写入 SMEM，TMA 引擎会更新进度；所有预期字节到达后，completion barrier 才会完成。 2. **Compute：** `tcgen05` MMA 从 SMEM 读取 operand tile，并把乘积累加到 TMEM tile 中。一个选定的 thread 提交这次 MMA；计算完成后，硬件会向对应的 barrier 发出完成信号。 3. **Epilogue：** warpgroup 把 TMEM accumulator 读回寄存器，将结果转换成输出 dtype，再写回 GMEM。这一步通常会先经过 SMEM staging，也可能使用 TMA store 完成最终写回。 diff --git a/zh/chapter_clc/index.md b/zh/chapter_clc/index.md index fbc8569a..3498af9e 100644 --- a/zh/chapter_clc/index.md +++ b/zh/chapter_clc/index.md @@ -13,7 +13,7 @@ 假设输出矩阵被划分成 100 个 tiles。最直接的做法是启动 100 个 CTAs，让第 0 个 CTA 计算 tile 0，第 1 个 CTA 计算 tile 1，以此类推。GPU 通常无法同时运行全部 100 个 CTAs，因此会先运行其中一部分；某个 CTA 结束并释放资源后，硬件再启动后续 CTA，直到所有 tiles 都处理完毕。 -传统的 fixed-number persistent kernel 使用另一种方式：它只启动一组长期运行的 CTAs 或 clusters，让每个 worker 在循环中连续计算多个 tiles。这样可以减少 CTA 启动和重复准备工作的开销，但也带来了新的调度问题：一个 worker 完成当前 tile 后，下一块 tile 从哪里来？ +传统的 persistent kernel 使用另一种方式：它只启动固定数量的长期运行 CTAs 或 clusters，让每个 worker 在循环中连续计算多个 tiles。这样可以减少 CTA 启动和重复准备工作的开销，但也带来了新的调度问题：一个 worker 完成当前 tile 后，下一块 tile 从哪里来？ 本章介绍 Blackwell 提供的 Cluster Launch Control（CLC）。CLC kernel 的 launch grid 仍然覆盖全部 output tiles，但运行中的 worker 可以取消尚未开始的 CTA 或 cluster launch，并接管它的 coordinate。这样既保留了完整 grid 的任务编号，又能让已经驻留的 workers 根据实际完成情况动态领取工作。 @@ -104,7 +104,7 @@ while true: 为什么要在计算当前 tile 之前请求下一块工作？因为 grid scheduler 处理请求需要时间。如果等当前 tile 算完才提交请求，这段延迟会直接落在两块 tile 之间，worker 只能停下来等待。 -提前提交后，scheduler 处理请求和当前 tile 的计算可以同时进行。等当前 tile 完成时，下一块工作的 coordinate 往往已经写入 shared memory。TMA 用计算覆盖数据搬运延迟，CLC 则用当前 tile 的计算覆盖调度请求延迟，两者采用的是同一种异步流水思路。 +提前提交后，scheduler 可以在当前 tile 计算期间处理请求。等当前 tile 完成时，下一块工作的 coordinate 往往已经写入 shared memory。TMA 将数据搬运延迟隐藏在计算之后，CLC 则用同样的异步流水思路隐藏调度请求的延迟。 CLC 通过 async proxy 把 response 写入 shared memory，普通 thread 则通过 generic proxy 查询这份结果。`mbarrier` wait 用来确认异步 response 已经写完；实际代码在提交新请求前和读完 response 后，还必须按照 PTX 要求执行相应的 proxy fence，建立 async proxy 与 generic proxy 之间的访问顺序，防止下一轮异步写入与尚未结束的读取发生冲突。此外，kernel 还需要正确处理 barrier phase，以及 CTA 或 cluster 范围的 thread synchronization。 diff --git a/zh/chapter_data_layout/index.md b/zh/chapter_data_layout/index.md index 22f2e96f..ac7b7e31 100644 --- a/zh/chapter_data_layout/index.md +++ b/zh/chapter_data_layout/index.md @@ -11,7 +11,7 @@ 同一组数字，如果以不同的物理排列方式写入内存，在同一块 GPU 上的运行速度可能相差一个数量级。 -原因在于，张量的逻辑索引并不说明它的字节在物理上实际存放在哪里。硬件对这种位置关系非常敏感：它决定 32 个 lane 的 load 能否合并成一次 transaction，还是分散成 32 次；决定这些地址会落到不同的 memory bank，还是撞到同一个 bank 并被串行化；甚至还决定一个 tile 的字节排列是否符合 Tensor Core 能够读取的格式。 +原因在于，张量的逻辑索引并不说明它的字节在物理上实际存放在哪里。硬件对这种位置关系非常敏感：根据访问模式，32 个 lanes 的 loads 可能合并成一次 transaction，也可能最多被拆成 32 次；这些地址可能落到不同的 memory banks，也可能集中到同一个 bank 并被串行处理；一个 tile 的字节排列还决定 Tensor Core 能否直接读取它。 机器学习程序通常用逻辑 shape 来描述张量。**数据布局**补上了缺失的物理部分：它说明带有逻辑索引 `(i, j, …)` 的元素实际存放在哪里，可以是在 memory 中、register 中，也可以是在其他硬件存储空间中。 @@ -405,7 +405,7 @@ bank = mapped_col 以第 1 行为例，交互图显示的逻辑列号依次为 `1, 0, 3, 2, 5, 4, 7, 6`。官方示意图在同一排列上加了这一行的编号偏移 `8`，因此显示为 `9, 8, 11, 10, 13, 12, 15, 14`。 -下面把图中的每个 128-bit cell 称为一个 16 B sector。对于 `SWIZZLE_128B`，atom 的每一行包含 8 个 sector，共 128 B。在常见的 4-byte bank 粒度下，一个 sector 横跨 4 个 bank，一整行正好覆盖 32 个 bank。swizzle 根据行坐标，对这一行中的 8 个 sector 做 XOR 重排。 +下面把图中的每个 128-bit cell 称为一个 16 B sector。对于 `SWIZZLE_128B`，atom 的每一行包含 8 个 sector，共 128 B。在常见的 4-byte bank 粒度下，一个 sector 中的四个 32-bit words 会映射到四个相邻 banks，一整行正好覆盖 32 个 banks。swizzle 根据行坐标，对这一行中的 8 个 sectors 做 XOR 重排。 一个 `SWIZZLE_128B` atom 包含 8 行，因此大小为 `8 × 128 B = 1024 B`。这里的 `128 B` 指 atom 每一行在连续维度上的宽度，而不是 atom 的总大小。atom 是地址重排的最小重复块，更大的 tile 由多个 atom 平铺而成。 diff --git a/zh/chapter_flash_attention/index.md b/zh/chapter_flash_attention/index.md index 666c3795..94bfd26b 100644 --- a/zh/chapter_flash_attention/index.md +++ b/zh/chapter_flash_attention/index.md @@ -33,13 +33,13 @@ $q_i$ 和 $k_j$ 的点积得到位置 $(i,j)$ 上的标量 score： $$s_{ij}=q_i\cdot k_j$$ -固定第 $i$ 个 query vector $q_i$ 后，让它分别与所有 key vectors $k_j$ 做点积，就得到这一行的 scores $s_{ij}$。这些 scores 组成 score matrix $S=QK^\top$ 的第 $i$ 行。基础 online softmax 先取这一行的最大值，作为计算指数时的参考值： +固定第 $i$ 个 query vector $q_i$ 后，让它分别与所有 key vectors $k_j$ 做点积，就得到这一行的 scores $s_{ij}$。这些 scores 组成 score matrix $S=QK^\top$ 的第 $i$ 行。把这一行真实的最大 score 记为 $m_i^{\max}$： -$$m_i=\max_j s_{ij}$$ +$$m_i^{\max}=\max_j s_{ij}$$ -计算指数前统一减去 $m_i$，可以让这一行最大的指数输入变成 0，避免指数值过大。这项平移会同时作用于 softmax 的分子和分母，因此不会改变最终的归一化结果。每个位置的未归一化 attention weight 为： +基础的稳定 softmax 会用 $m_i^{\max}$ 作为指数参考值。计算指数前统一减去它，可以让这一行最大的指数输入变成 0，避免指数值过大。这项平移会同时作用于 softmax 的分子和分母，因此不会改变最终的归一化结果。每个位置的未归一化 attention weight 为： -$$p_{ij}=\exp\left(\frac{s_{ij}-m_i}{\sqrt d}\right)$$ +$$p_{ij}=\exp\left(\frac{s_{ij}-m_i^{\max}}{\sqrt d}\right)$$ 将这一行的所有 $p_{ij}$ 相加，得到未归一化权重之和 $\ell_i$。再用同一组 $p_{ij}$ 对 value vectors 加权求和，得到尚未除以 $\ell_i$ 的 output vector $o_i$： @@ -51,7 +51,7 @@ $$o_i=\sum_j p_{ij}v_j$$ $$O_i=\frac{o_i}{\ell_i}$$ -FlashAttention 按 block 处理 K/V。一个 block 的 scores 使用完后就可以丢弃，后续计算只需要保留当前的 $m_i$、$\ell_i$ 和 $o_i$。其中，$\ell_i$ 和 $o_i$ 都是使用当时的参考值 $m_i$ 计算并累积的。因此，后续 block 一旦改用更大的参考值，旧状态就必须先换算到新尺度，才能与当前 block 的贡献相加。 +FlashAttention 按 block 处理 K/V。一个 block 的 scores 使用完后就可以丢弃；kernel 只需为每一行保留指数参考值 $r_i$、running denominator $\ell_i$ 和 running weighted sum $o_i$。基础 online softmax 会把 $r_i$ 更新为截至当前最大的 score，而 FA4 可以暂时保留旧值。$\ell_i$ 和 $o_i$ 都是相对于当前 $r_i$ 累加的，因此后续 block 一旦改用更大的参考值，旧状态就必须先换算到新尺度，才能与当前 block 的贡献相加。 基础 online softmax 每次发现更大的逐行最大值都会完成这次换算。FA4 则先比较新旧参考值的差距：差距较小时继续使用旧值，从而避免立即重缩放已经累积的 output。要理解这项优化，先把参考值变化时的尺度转换写清楚。 @@ -63,35 +63,35 @@ $$\alpha=\frac{\log_2(e)}{\sqrt d}$$ $$\exp\left(\frac{s-m}{\sqrt d}\right)=2^{(s-m)\alpha}$$ -代码将 $\alpha$ 记为 `scale_log2`。设旧状态使用参考值 $m_{\mathrm{old}}$，当前 block 的逐行最大值为 $m_{\mathrm{block}}$。用下标 $c$ 表示 candidate，本轮可选的新参考值为： +代码将 $\alpha$ 记为 `scale_log2`。设旧状态使用参考值 $r_{\mathrm{old}}$，当前 block 的逐行最大值为 $m_{\mathrm{block}}$。用下标 $c$ 表示 candidate，本轮可选的新参考值为： -$$m_c=\max(m_{\mathrm{old}},m_{\mathrm{block}})$$ +$$r_c=\max(r_{\mathrm{old}},m_{\mathrm{block}})$$ 再定义二者在 base-2 exponent 中的有符号差距 $\delta$，它对应代码变量 `delta`： -$$\delta=(m_{\mathrm{old}}-m_c)\alpha\le 0$$ +$$\delta=(r_{\mathrm{old}}-r_c)\alpha\le 0$$ -$\delta$ 是旧参考值减去候选参考值后的有符号结果；$-\delta$ 才表示候选参考值高出了多少个 base-2 exponent units。由于 $m_c\ge m_{\mathrm{old}}$，$\delta$ 不会大于 0。 +$\delta$ 是旧参考值减去候选参考值后的有符号结果；$-\delta$ 才表示候选参考值高出了多少个 base-2 exponent units。由于 $r_c\ge r_{\mathrm{old}}$，$\delta$ 不会大于 0。 在 [FA4 论文](https://arxiv.org/abs/2603.05451)中，阈值通常取 $\tau=\log_2(256)=8$。当 $-\delta=8$ 时，继续使用旧参考值会让当前 block 的最大未归一化权重达到 $2^8=256$；若切换到候选参考值，旧状态则要乘 $2^\delta=1/256$。因此，阈值 8 表示允许新旧尺度相差最多 256 倍，超过后才执行重缩放：`delta >= -8` 时保留旧参考值，`delta < -8` 时切换参考值。这种通过阈值延迟重缩放的做法，是 FA4 为减少 correction 的数据搬运和乘法开销而引入的执行优化；取值 8 则在减少重缩放次数和限制指数增长之间作了折中。 -如果本轮改用候选参考值 $m_c$，此前相对于旧参考值计算的每个指数都要乘同一个系数： +如果本轮改用候选参考值 $r_c$，此前相对于旧参考值计算的每个指数都要乘同一个系数： -$$e^{(s-m_c)/\sqrt d} -=e^{(s-m_{\mathrm{old}})/\sqrt d} -\cdot e^{(m_{\mathrm{old}}-m_c)/\sqrt d}$$ +$$e^{(s-r_c)/\sqrt d} +=e^{(s-r_{\mathrm{old}})/\sqrt d} +\cdot e^{(r_{\mathrm{old}}-r_c)/\sqrt d}$$ 将这个尺度转换系数记为 $a_{\mathrm{scale}}$，则： $$a_{\mathrm{scale}} -=e^{(m_{\mathrm{old}}-m_c)/\sqrt d} +=e^{(r_{\mathrm{old}}-r_c)/\sqrt d} =2^\delta$$ -切换到候选参考值 $m_c$ 后，之前累积的归一化分母 $\ell_i$ 和未归一化加权和 $o_i$ 仍处于旧尺度。Kernel 先将两者同时乘以 $a_{\mathrm{scale}}=2^\delta$，转换到新尺度，再与当前 block 的结果相加。下面映射到伪代码时，$\ell_i$ 和 $o_i$ 分别记为 `row_sum` 和 `O`，转换系数则由 `acc_scale = exp2(delta)` 计算。 +切换到候选参考值 $r_c$ 后，之前累积的归一化分母 $\ell_i$ 和未归一化加权和 $o_i$ 仍处于旧尺度。Kernel 先将两者同时乘以 $a_{\mathrm{scale}}=2^\delta$，转换到新尺度，再与当前 block 的结果相加。下面映射到伪代码时，$\ell_i$ 和 $o_i$ 分别记为 `row_sum` 和 `O`，转换系数则由 `acc_scale = exp2(delta)` 计算。 对应到下面的伪代码，需要跨 K/V blocks 保留的三项状态分别是： -- `row_max`：计算指数时从这一行所有 scores 中共同减去的参考值，也就是 $m_i$。基础算法使用当前最大 score；FA4 在阈值允许时可以继续使用旧值。 +- `row_max`：计算指数时从这一行所有 scores 中减去的参考值 $r_i$。基础 online softmax 使用截至当前最大的 score；FA4 在阈值允许时可以继续使用旧参考值。因此，尽管变量名是 `row_max`，它并不保证在每个 iteration 都等于真实最大值 $m_i^{\max}$。 - `row_sum`：已经处理过的所有 key positions 的 $p_{ij}$ 之和，也就是 $\ell_i$。 - `O`：使用同一组 $p_{ij}$ 得到的加权和 $o_i$；所有 blocks 处理完成后再除以 `row_sum`。 @@ -721,7 +721,7 @@ if any_needs_rescale != 0: # 当前 warp：TMEM -> registers -> multiply -> TMEM ... -# correction loop 交错归还另一个 Q stage +# correction loop 在此归还另一个 Q stage p_o_rescale.arrive(i_q) softmax_corr.empty.arrive(1 - i_q) ``` @@ -827,7 +827,7 @@ while scheduler.valid(): m_block_idx = scheduler.m_block_idx batch_idx = scheduler.batch_idx kv_head_idx = scheduler.head_idx - # process one Q block against its K/V block range + # 使用对应范围内的 K/V blocks 处理一个 Q block scheduler.next_tile() ``` @@ -872,7 +872,7 @@ with target: ex.mod(Q, K, V, O, prof) torch.cuda.synchronize() -# torch reference; enable_gqa lets the 32 query heads share the 8 KV heads +# torch reference；enable_gqa 允许 32 个 query heads 共享 8 个 KV heads qt, kt, vt = (x.transpose(1, 2).float() for x in (Q, K, V)) ref = F.scaled_dot_product_attention(qt, kt, vt, enable_gqa=True).transpose(1, 2).half() torch.testing.assert_close(O, ref, rtol=1e-2, atol=1e-2) diff --git a/zh/chapter_gemm_advanced/index.md b/zh/chapter_gemm_advanced/index.md index 898a9451..2f8bea4b 100644 --- a/zh/chapter_gemm_advanced/index.md +++ b/zh/chapter_gemm_advanced/index.md @@ -57,7 +57,7 @@ Load 与 MMA 之间通过两个 barriers 交接 SMEM buffer： ### 四个 Barriers -三个并发角色之间需要四个 barriers。正向路径 TMA → MMA → Writeback 表示数据已经准备好；反向路径 Writeback → MMA → TMA 表示 buffer 已经释放。Barrier 名称采用 `source2destination`，例如 `tma2mma` 表示 TMA 向 MMA 发送通知。 +三个并发角色之间需要四个 barriers。正向路径 TMA → MMA → Writeback 表示数据已经准备好；反向路径 Writeback → MMA → TMA 则把各自保护的 buffer 或资源交还给前一角色复用。Barrier 名称采用 `source2destination`，例如 `tma2mma` 表示 TMA 向 MMA 发送通知。 | Barrier | 类型 | 方向 | 含义 | |---------|------|------|------| @@ -76,8 +76,8 @@ Barrier 类型取决于 producer 如何报告完成。**TMA load** 使用带 byt ```python tma_ps = PipelineState(PIPE_DEPTH, phase=1) # Producer starts ready (phase=1) -# tma_ps.stage = current stage index -# tma_ps.phase = current phase (0 or 1) +# tma_ps.stage 表示当前 stage index +# tma_ps.phase 表示当前 phase（0 或 1） tma_ps.advance() # Advance to next stage ``` @@ -200,7 +200,7 @@ def hgemm_v7(M, N, K): n_st = T.meta_var(tile_scheduler.n_idx * BLK_N) # ============================================= - # Warpgroup 1: TMA Producer (warp 3) + MMA Consumer (warp 0) + # Warpgroup 1：TMA producer（warp 3）+ MMA consumer（warp 0） # ============================================= if wg_id == 1: if warp_id == 3: @@ -235,7 +235,7 @@ def hgemm_v7(M, N, K): if T.filter(lane_id, T.ptx.elect_sync()): while tile_scheduler.valid(): - # Wait for TMEM to be free from previous tile's writeback + # 等待上一块 tile 的 writeback 释放 TMEM ld2mma.wait(ld_ps.stage, ld_ps.phase) ld_ps.advance() @@ -249,37 +249,37 @@ def hgemm_v7(M, N, K): mma2tma.arrive(mma_ps.stage, cta_group=1, cta_mask=0) mma_ps.advance() - # Signal results ready for writeback + # 通知 writeback：结果已经准备好 mma2ld.arrive(0, cta_group=1, cta_mask=0) tile_scheduler.next_tile() # ============================================= - # Warpgroup 0: Writeback + # Warpgroup 0：writeback # ============================================= elif wg_id == 0: wb_ps = PipelineState(1, phase=0) reg_f16 = T.alloc_local((BLK_N,), d_type) while tile_scheduler.valid(): - # Wait for MMA results + # 等待 MMA 结果 mma2ld.wait(wb_ps.stage, wb_ps.phase) wb_ps.advance() T.ptx.tcgen05.fence.after_thread_sync() - # Read TMEM -> registers (warpgroup scope) + # 以 warpgroup scope 读取 TMEM -> registers reg = T.alloc_local((BLK_N,), acc_type) reg_wg = reg.view(128, BLK_N, layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) Tx.wg.copy_async(reg_wg[:], tmem[:, :BLK_N]) T.ptx.tcgen05.wait.ld() - # Signal TMEM free (all 128 threads arrive) + # 所有 128 个 threads 报告 arrival，通知 MMA 可以复用 TMEM ld2mma.arrive(0, cta_id=0, pred=True) - # Cast fp32 -> fp16 + # 转换 fp32 -> fp16 Tx.cast(reg_f16[:], reg[:]) - # Write to Dsmem + TMA store + # 写入 Dsmem，再执行 TMA store Tx.copy(Dsmem[warp_id * 32 + lane_id, :], reg_f16[:]) T.ptx.fence.proxy_async("shared::cta") T.cuda.warpgroup_sync(10) @@ -500,7 +500,7 @@ def hgemm_v8(M, N, K): tma2mma_cta0 = tma2mma.remote_view(0) # ============================================= - # Warpgroup 1: TMA Producer (warp 3) + MMA Consumer (warp 0) + # Warpgroup 1：TMA producer（warp 3）+ MMA consumer（warp 0） # ============================================= if wg_id == 1: if warp_id == 3: @@ -552,7 +552,7 @@ def hgemm_v8(M, N, K): tile_scheduler.next_tile() # ============================================= - # Warpgroup 0: Writeback (256 columns in 2 x 128-column chunks) + # Warpgroup 0：writeback（将 256 columns 分成两个 128-column chunks） # ============================================= elif wg_id == 0: wb_ps = PipelineState(1, phase=0) @@ -750,7 +750,7 @@ def hgemm_v9(M, N, K): tma2mma_cta0 = tma2mma.remote_view(0) # ============================================= - # Warpgroup 2: TMA Producer (warp 3) + 2 MMA Consumers (warp 0, 1) + # Warpgroup 2：TMA producer（warp 3）+ 两个 MMA consumers（warp 0、1） # ============================================= if wg_id == 2: if warp_id == 3: @@ -809,7 +809,7 @@ def hgemm_v9(M, N, K): tile_scheduler.next_tile() # ============================================= - # Warpgroup 0/1: Writeback (each reads its consumer's TMEM range) + # Warpgroup 0/1：writeback（分别读取对应 consumer 的 TMEM range） # ============================================= elif wg_id < NUM_CONSUMER: wb_ps = PipelineState(1, phase=0) @@ -820,7 +820,7 @@ def hgemm_v9(M, N, K): wb_ps.advance() T.ptx.tcgen05.fence.after_thread_sync() - # Read TMEM in EPI_N=64 column chunks (4 iterations for 256 cols) + # 以 EPI_N=64 为单位分块读取 TMEM（256 columns 共需四轮） for i in T.unroll(MMA_N // EPI_N): reg = T.alloc_local((EPI_N,), acc_type) reg_wg = reg.view(128, EPI_N, @@ -876,7 +876,7 @@ def hgemm_v9(M, N, K): 表中给出具体时间的版本都在相同的 `M=N=K=4096` 规模下测量，因此可以直接比较。第 1 步的 70 ms 来自一个采用相同串行数据路径的完整矩阵 baseline，并不是直接运行 {ref}`chap_gemm_basics` 中只计算一个 $128\times128$ tile 的 `hgemm_v1`。基础章节使用较小规模讲解第 1 至 3 步；表中的第 1、3 步则是相应思路扩展到完整矩阵后的测量结果。 -第 2 步仍然只计算一个 output tile，不能与表中的完整矩阵结果直接比较。第 5、6 步则是从 TMA load 逐步过渡到 warp specialization 的中间版本，相关机制都包含在第 7 步中；表格只保留这一段的起点和终点。因此，第 2、5、6 步以横线表示，也不计算对应的单步加速比。 +第 2 步仍然只计算一个 output tile，不能与表中的完整矩阵结果直接比较。第 5、6 步则是从 TMA load 逐步过渡到 warp specialization 的中间版本，相关机制都包含在第 7 步中；表格只保留这一段的起点和终点。因此，第 2、5、6 步以横线表示，不展示它们相对第 1 步的累计加速比。 这些数字来自同一次 B200 reference run，只用于比较本章各版本在相同条件下的相对变化，不代表其他输入规模或测试环境下的硬件峰值。 diff --git a/zh/chapter_gemm_async/index.md b/zh/chapter_gemm_async/index.md index 1de55ab5..f389db58 100644 --- a/zh/chapter_gemm_async/index.md +++ b/zh/chapter_gemm_async/index.md @@ -170,41 +170,41 @@ def hgemm_v4(M, N, K): for k in range(K_TILES): k_st = T.meta_var(k * BLK_K) - # Single thread issues TMA load + # 由一个 thread 发起 TMA load if tid == 0: tma_load(k_st) - # Wait for TMA to finish; the mbarrier release carries SMEM - # visibility to the subsequent MMA, so no extra fence is needed. + # 等待 TMA 完成；mbarrier 提供后续 MMA 读取 SMEM 所需的可见性， + # 因此这里不需要额外的 fence。 T.ptx.mbarrier.try_wait(tma_bar.ptr_to([0]), phase_tma) - # Single thread issues MMA + # 由一个 thread 发起 MMA if tid == 0: mma(accum=k != 0) - # Wait for MMA to finish + # 等待 MMA 完成 T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) phase_tma ^= 1 phase_mma ^= 1 - # --- TMA Store Writeback --- + # --- 使用 TMA store 写回 --- Dreg = T.alloc_local((BLK_N,), acc_type) Dreg_f16 = T.alloc_local((BLK_N,), d_type) Dreg_wg = Dreg.view(128, BLK_N, layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) - # Read TMEM -> registers (async; wait.ld then cta_sync to ensure read completes) + # 异步读取 TMEM -> registers；先执行 wait.ld，再用 cta_sync 同步 threads Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) T.ptx.tcgen05.wait.ld() T.cuda.cta_sync() - # Cast fp32 -> fp16 + # 转换 fp32 -> fp16 Tx.cast(Dreg_f16[:], Dreg[:]) - # Write registers -> Dsmem, flush, then sync + # 写入 registers -> Dsmem，建立可见性后再同步 Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:]) T.ptx.fence.proxy_async("shared::cta") T.cuda.warpgroup_sync(10) - # TMA store: Dsmem -> GMEM. One selected thread starts the store and drains the - # store group before Dsmem is reused. + # TMA store：Dsmem -> GMEM。一个 selected thread 发起 store； + # 复用 Dsmem 前必须等待该 store group 完成。 if tid == 0: Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N], Dsmem[:, :], dispatch="tma") @@ -212,7 +212,7 @@ def hgemm_v4(M, N, K): T.ptx.cp_async.bulk.wait_group(0) T.cuda.warpgroup_sync(10) - # --- Deallocate TMEM --- + # --- 释放 TMEM --- T.cuda.cta_sync() if warp_id == 0: T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) @@ -316,7 +316,7 @@ def hgemm_v5(M, N, K): BLK_M, BLK_N, BLK_K = 128, 128, 64 K_TILES = K // BLK_K - # Double-buffered layouts: first dimension is pipeline stage + # 双缓冲 layout：第一维表示 pipeline stage A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_M, BLK_K)) B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, @@ -339,7 +339,7 @@ def hgemm_v5(M, N, K): # --- SMEM allocation --- pool = T.SMEMPool() tmem_addr = pool.alloc((1,), "uint32") - # Double-buffered TMA barriers (one per stage), single MMA barrier + # 每个双缓冲 stage 使用一个 TMA barrier；所有 stages 共用一个 MMA barrier tma_bar = pool.alloc((PIPE_DEPTH,), "uint64", align=8) mma_bar = pool.alloc((1,), "uint64", align=8) pool.move_base_to(1024) @@ -348,7 +348,7 @@ def hgemm_v5(M, N, K): Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout) pool.commit() - # Initialize barriers: PIPE_DEPTH for TMA, 1 for MMA + # 初始化 barriers：TMA 使用 PIPE_DEPTH 个，MMA 使用 1 个 if warp_id == 0: if lane_id == 0: T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) @@ -404,23 +404,23 @@ def hgemm_v5(M, N, K): for k in range(K_TILES): stage = k % PIPE_DEPTH - # Wait for TMA to finish loading this stage + # 等待 TMA 完成当前 stage 的加载 T.ptx.mbarrier.try_wait(tma_bar.ptr_to([stage]), phase_tma) - # MMA on this stage's data + # 使用当前 stage 的数据执行 MMA if tid == 0: mma(stage, accum=(k != 0)) T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) phase_mma ^= 1 - # Issue next prefetch load (k + PIPE_DEPTH) + # 发起下一次 prefetch（k + PIPE_DEPTH） next_k = k + PIPE_DEPTH if next_k < K_TILES: if tid == 0: tma_load(stage, next_k * BLK_K) - # TMA phase flips when stage wraps around + # stage index 绕回时翻转 TMA phase if stage == PIPE_DEPTH - 1: phase_tma ^= 1 @@ -443,7 +443,7 @@ def hgemm_v5(M, N, K): T.ptx.cp_async.bulk.wait_group(0) T.cuda.warpgroup_sync(10) - # Deallocate TMEM + # 释放 TMEM T.cuda.cta_sync() if warp_id == 0: T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) @@ -579,7 +579,7 @@ def hgemm_v6(M, N, K): layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]) ) - # Tile scheduler: assigns tiles to CTAs in L2-friendly order + # Tile scheduler：按有利于 L2 locality 的顺序将 tiles 分配给 CTAs tile_scheduler = ClusterPersistentScheduler2D( "ts", num_m_tiles=M // BLK_M, @@ -615,7 +615,7 @@ def hgemm_v6(M, N, K): # === Outer loop: iterate over tiles === while tile_scheduler.valid(): - # Get current tile position from scheduler + # 从 scheduler 取得当前 tile 坐标 m_st = T.meta_var(tile_scheduler.m_idx * BLK_M) n_st = T.meta_var(tile_scheduler.n_idx * BLK_N) @@ -623,7 +623,7 @@ def hgemm_v6(M, N, K): phase_tma: T.int32 = 0 phase_mma: T.int32 = 0 - # Prefetch first PIPE_DEPTH stages + # 预取最初的 PIPE_DEPTH 个 stages if tid == 0: for s in range(min(PIPE_DEPTH, K_TILES)): tma_load(s, s * BLK_K, m_st, n_st) @@ -665,7 +665,7 @@ def hgemm_v6(M, N, K): T.cuda.cta_sync() tile_scheduler.next_tile() # Move to next tile - # Deallocate TMEM + # 释放 TMEM T.cuda.cta_sync() if warp_id == 0: T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) diff --git a/zh/chapter_gemm_basics/index.md b/zh/chapter_gemm_basics/index.md index 46ef5476..50ff04f9 100644 --- a/zh/chapter_gemm_basics/index.md +++ b/zh/chapter_gemm_basics/index.md @@ -70,7 +70,7 @@ $$\text{TFLOPS} = \frac{2 \times M \times N \times K}{t_{\text{seconds}} \times 这个 kernel 只沿 `GMEM -> SMEM -> TMEM -> registers -> GMEM` 路径执行一次，不包含循环。具体步骤如下： 1. **分配**：通过 pool allocator 分配 SMEM，通过 `tcgen05.alloc` 分配 TMEM，并准备等待 MMA 完成的 mbarrier。 -2. **加载**：128 个 threads 使用同步 `Tx.copy`，协作将 A、B tiles 从 GMEM 搬到 SMEM。 +2. **加载**：128 个 threads 使用同步 `Tx.cta.copy`，协作将 A、B tiles 从 GMEM 搬到 SMEM。 3. **计算**：选出的一个 thread 发出 `Tx.gemm_async` 和 `tcgen05.commit`，所有 threads 等待 mbarrier。 4. **写回**：warpgroup 将 TMEM 读入 registers；每个 thread 把 fp32 转成 fp16，再写入 GMEM。 5. **释放**：释放 TMEM。 @@ -184,9 +184,8 @@ def hgemm_v1(M, N, K): acc_type = tvm.DataType("float32") BLK_M, BLK_N, BLK_K = 128, 128, 64 - # MMA_M/MMA_N/MMA_K document the underlying hardware MMA tile; they are not - # passed to gemm_async (which derives the MMA shape from the operand and - # accumulator tiles), so the later steps omit them. + # MMA_M/MMA_N/MMA_K 记录底层硬件 MMA tile 的 shape。gemm_async 会根据 + # operands 和 accumulator tiles 推导该 shape，因此后续步骤不再保留这些常量。 MMA_M, MMA_N, MMA_K = 128, 128, 16 A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K)) @@ -199,9 +198,8 @@ def hgemm_v1(M, N, K): D: T.Buffer((M, N), d_type), ): T.device_entry() - # Step 1 is a single-tile kernel: M = BLK_M and N = BLK_N, so the grid - # is 1x1. Starting with a 1x1 grid keeps the per-CTA tile offsets - # (m_st, n_st) trivially zero; Steps 3+ generalise this to larger M / N. + # 第 1 步只计算一个 tile：M=BLK_M、N=BLK_N，因此 grid shape 为 1x1。 + # 此时每个 CTA 的 tile offsets（m_st、n_st）都为 0；第 3 步再扩展到更大的 M、N。 bx, by = T.cta_id([M // BLK_M, N // BLK_N]) wg_id = T.warpgroup_id([1]) # single warpgroup, so wg_id is always 0 (unused below) warp_id = T.warp_id_in_wg([4]) @@ -235,14 +233,14 @@ def hgemm_v1(M, N, K): n_st = T.meta_var(by * BLK_N) phase_mma: T.int32 = 0 - # --- Load: all threads copy global -> shared (synchronous). - # With M=BLK_M and N=BLK_N the slices below cover the full matrices; - # the slice form is kept so the diff to Step 3 (multi-tile) is minimal. + # --- Load：所有 threads 同步完成 global -> shared copy --- + # M=BLK_M、N=BLK_N 时，下面的 slices 覆盖完整矩阵；保留 slice 写法， + # 便于与第 3 步的 multi-tile 版本比较。 Tx.cta.copy(Asmem[:, :], A[m_st:m_st + BLK_M, :]) Tx.cta.copy(Bsmem[:, :], B[n_st:n_st + BLK_N, :]) T.cuda.cta_sync() - # --- Compute: single elected thread issues MMA --- + # --- Compute：由一个 elected thread 发起 MMA --- if warp_id == 0: if T.ptx.elect_sync(): Tx.gemm_async( @@ -253,7 +251,7 @@ def hgemm_v1(M, N, K): T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) - # --- Writeback: TMEM -> RF -> GMEM --- + # --- Writeback：TMEM -> RF -> GMEM --- Dreg = T.alloc_local((BLK_N,), acc_type) Dreg_f16 = T.alloc_local((BLK_N,), d_type) Dreg_wg = Dreg.view(128, BLK_N, @@ -264,7 +262,7 @@ def hgemm_v1(M, N, K): m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:]) - # --- Deallocate TMEM --- + # --- 释放 TMEM --- T.cuda.cta_sync() if warp_id == 0: T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) @@ -292,18 +290,18 @@ A_tensor = torch.randn(M, K, dtype=torch.float16, device=device) B_tensor = torch.randn(N, K, dtype=torch.float16, device=device) D_tensor = torch.zeros(M, N, dtype=torch.float16, device=device) -# ex.mod(...) takes torch tensors directly, the same call form used in every chapter. +# ex.mod(...) 可以直接接收 torch tensors；后续章节沿用相同的调用方式。 ex.mod(A_tensor, B_tensor, D_tensor) D_ref = (A_tensor.float() @ B_tensor.float().T).half() max_err = float((D_tensor - D_ref).abs().max()) print(f"Max error vs torch reference: {max_err:.6f}") -# Relative tolerance, like the warp-specialization and Flash Attention cells: -# output magnitude grows with K, so a fixed absolute bound would fail at larger K. +# 与 warp specialization 和 Flash Attention 的示例一样，这里使用相对容差： +# output magnitude 会随 K 增长，固定的绝对误差上限不适用于较大的 K。 torch.testing.assert_close(D_tensor, D_ref, rtol=2e-2, atol=1e-2) print("PASS") -# Optional timing for larger kernels. +# 对更大 kernel 进行可选计时。 ITERS = 10 for _ in range(3): ex.mod(A_tensor, B_tensor, D_tensor) @@ -404,7 +402,7 @@ def hgemm_v2(M, N, K): D: T.Buffer((M, N), d_type), ): T.device_entry() - bx, by = T.cta_id([M // BLK_M, N // BLK_N]) # still one output tile (M=N=128) + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) # 仍然只有一个 output tile（M=N=128） wg_id = T.warpgroup_id([1]) warp_id = T.warp_id_in_wg([4]) lane_id = T.lane_id([32]) @@ -434,26 +432,26 @@ def hgemm_v2(M, N, K): m_st = T.meta_var(bx * BLK_M) n_st = T.meta_var(by * BLK_N) - # === K-loop: iterate over K in chunks of BLK_K === - for i in T.serial(K_TILES): # serial device loop (keeps the full-K A/B parameters correctly shaped) - # Load the i-th K chunk + # === K-loop：以 BLK_K 为单位遍历 K === + for i in T.serial(K_TILES): # device 侧串行 loop；A、B parameters 仍保留完整 K 维 + # 加载第 i 个 K chunk Tx.cta.copy(Asmem[:, :], A[:, i*BLK_K:(i+1)*BLK_K]) Tx.cta.copy(Bsmem[:, :], B[:, i*BLK_K:(i+1)*BLK_K]) T.cuda.cta_sync() - # MMA: accum=False for first tile, True for rest + # 第一个 tile 使用 accum=False，后续 tiles 使用 accum=True if warp_id == 0: if T.ptx.elect_sync(): Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], accum=(i != 0), dispatch="tcgen05", cta_group=1) T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) - # Wait for MMA, then flip phase + # 等待 MMA 完成，再翻转 phase T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) phase_mma ^= 1 - # === Writeback (same as Step 1) === + # === Writeback（与第 1 步相同）=== Dreg = T.alloc_local((BLK_N,), acc_type) Dreg_f16 = T.alloc_local((BLK_N,), d_type) Dreg_wg = Dreg.view(128, BLK_N, @@ -585,12 +583,12 @@ def hgemm_v3(M, N, K): phase_mma: T.int32 = 0 - # Per-CTA tile offsets + # 当前 CTA 的 tile offsets m_st = T.meta_var(bx * BLK_M) n_st = T.meta_var(by * BLK_N) - # K-loop with offset A and B slices - for i in T.serial(K_TILES): # serial device loop (keeps the full-K A/B parameters correctly shaped) + # K-loop：加载带 offset 的 A、B slices + for i in T.serial(K_TILES): # device 侧串行 loop；A、B parameters 仍保留完整 K 维 Tx.cta.copy(Asmem[:, :], A[m_st:m_st+BLK_M, i*BLK_K:(i+1)*BLK_K]) Tx.cta.copy(Bsmem[:, :], B[n_st:n_st+BLK_N, i*BLK_K:(i+1)*BLK_K]) @@ -605,7 +603,7 @@ def hgemm_v3(M, N, K): T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) phase_mma ^= 1 - # Writeback to the correct output tile + # 写回当前 CTA 对应的 output tile Dreg = T.alloc_local((BLK_N,), acc_type) Dreg_f16 = T.alloc_local((BLK_N,), d_type) Dreg_wg = Dreg.view(128, BLK_N, @@ -628,6 +626,6 @@ def hgemm_v3(M, N, K): ## 练习 -1. 在第 1 至第 3 步中，`Tx.copy` 会在 MMA 之前将 A、B tiles 搬入 SMEM。为什么 `Tx.gemm_async` 读取这些 tiles 前必须执行 `T.cuda.cta_sync()`？ +1. 在第 1 至第 3 步中，`Tx.cta.copy` 会在 MMA 之前将 A、B tiles 搬入 SMEM。为什么 `Tx.gemm_async` 读取这些 tiles 前必须执行 `T.cuda.cta_sync()`？ 2. 在第 2 步中，如果从 K-loop 删除 `phase_mma ^= 1`，会发生什么？Kernel 仍会等待每次 MMA，还是后续 wait 可能提前通过？ 3. 当 `M=N=4096`、`BLK_M=BLK_N=128` 时，第 3 步的 grid shape 是多少，共启动多少个 CTAs？对于 CTA `(bx, by)`，哪些 CTAs 会独立读取相同的 A tiles，哪些会独立读取相同的 B tiles？当前 kernel 是否显式共享了这些数据？ diff --git a/zh/chapter_intro_tirx/index.md b/zh/chapter_intro_tirx/index.md index f044ce19..8b13af01 100644 --- a/zh/chapter_intro_tirx/index.md +++ b/zh/chapter_intro_tirx/index.md @@ -65,7 +65,7 @@ D: tcgen05.mma -> TMEM -> registers -> GMEM 3. 通过 `Tx.gemm_async` 发起 MMA； 4. 将结果从 TMEM 读回 registers，再写入 GMEM。 -其中最关键的三项 tile 操作是 `Tx.cta.copy`、`Tx.gemm_async` 和 `Tx.wg.copy_async`。其余 PTX 调用用于申请和释放 TMEM、初始化 barrier 并建立同步；本章先把它们看作完成这几个阶段所需的底层步骤。 +其中最关键的三项 tile 操作是 `Tx.cta.copy`、`Tx.gemm_async` 和 `Tx.wg.copy_async`。其余底层调用用于申请和释放 TMEM、初始化 barrier 并建立同步；本章先把它们看作完成这几个阶段所需的实现步骤。 先导入这个 kernel 使用的模块： diff --git a/zh/chapter_layout_generations/index.md b/zh/chapter_layout_generations/index.md index c7d30ba1..0e986e54 100644 --- a/zh/chapter_layout_generations/index.md +++ b/zh/chapter_layout_generations/index.md @@ -266,7 +266,7 @@ S[(4, 32, 4) : (4@TCol, 1@TLane, 1@TCol)] 其中，`S[...]` 把 `(Mgroup, lane, sfk)` 映射到 TMEM 中的 byte 位置。在带有数据类型的 TIRx layout 中，`@TCol` stride 以 buffer element 为单位。这里每个 scale factor 占 8 bits，因此逻辑 TCol 位置为 `4*Mgroup+sfk`；每四个连续位置打包进一个 32-bit hardware TCol cell。等价地，`hardware_TCol=(4*Mgroup+sfk)//4`，`byte_in_word=(4*Mgroup+sfk)%4`。 -`R[...]` 表示沿 `TLane` 轴复制四份。`tcgen05.cp` 的 `.32x128b.warpx4` 形式正好完成这件事：先写入一个 32-lane window，再把同一份数据广播到另外三个 warp windows。 +`R[...]` 表示沿 `TLane` 轴复制四份。`tcgen05.cp` 的 `.32x128b.warpx4` 形式会把同一个基础 tile multicast 到四个 32-lane warp windows，从而得到这一布局。 ### `scale_vec` 的 Word 内复制 diff --git a/zh/chapter_performance/index.md b/zh/chapter_performance/index.md index ba22f2ae..7011748e 100644 --- a/zh/chapter_performance/index.md +++ b/zh/chapter_performance/index.md @@ -189,7 +189,7 @@ Roofline 模型可以判断一个 kernel 的性能上限，但不会告诉我们 一个大规模 fp16 GEMM 在理论上可能是 compute-bound 的。但这只说明 HBM 层的内存上限不是主要瓶颈，并不意味着任意一种实现都能达到 Tensor Core 的计算上限。要缩小这中间的差距，需要正确的指令、layout、staging、同步和调度。后续 GEMM 章节会在 B200 上通过一系列步骤展示这一点：每一步都保持相同的基本算法，但改变 tile 的计算方式或调度方式。 -在 GEMM 的优化阶梯中，第一个明显的实测性能跃升，是从 thread-copy tiled 路径切换到 TMA-backed 路径。前者由 CTA 中的普通线程执行 GMEM 到 SMEM 的拷贝；后者把这种规则的 tile 搬运交给 TMA 硬件引擎，让 kernel 可以通过硬件管理的大块拷贝来持续为 Tensor Cores 提供数据。 +在 GEMM 的优化阶梯中，第一个明显的实测性能跃升，是从 thread-copy tiled 路径切换到 TMA-backed 路径。前者由 CTA 中的普通 threads 将 tiles 从 GMEM 搬到 SMEM；后者把这种规则的 tile 搬运交给 TMA 硬件引擎。TMA 通过硬件管理的大块 copy 填充 SMEM，MMA 随后再从 SMEM 读取这些 tiles。 在第一次跃升之后，后续优化都围绕一个问题展开：如何减少数据搬运、Tensor Core 计算和 epilogue 之间的等待。Software pipelining 和 warp specialization 会重新安排这些阶段，使不同硬件单元能够重叠工作。下一节具体说明这种调度方式。 @@ -201,7 +201,7 @@ Roofline 模型可以判断一个 kernel 的性能上限，但不会告诉我们 ## 通过重叠执行减少硬件空闲 -一旦 GEMM 已经是 compute-bound 的，并且已经使用了 Tensor Cores，剩下的性能差距通常来自于硬件的空闲时间。 +当 GEMM 已经进入 compute-bound 区域并使用 Tensor Cores 后，剩余的性能差距通常来自某些执行路径没有得到充分利用。 一个简单的 kernel 可能会这样执行： @@ -224,7 +224,7 @@ compute tile k store tile k - 1 ``` -在 Blackwell 上，这三个阶段分别主要由 TMA、`tcgen05.mma` 和 epilogue/store 路径完成，`mbarrier` 负责它们之间的数据交接。 +在 Blackwell 上，这三个阶段分别主要由 TMA、`tcgen05.mma` 和 epilogue/store 路径完成，`mbarrier` 则协调阶段完成状态和 buffer 的复用时机。 重叠执行并不会消除依赖关系：tile `k` 的 MMA 仍然必须等待它加载完成，epilogue 也必须等待 MMA 完成。可以提前执行的是与当前计算没有直接依赖的工作，例如加载 tile `k+1`，或者写回 tile `k-1`。 @@ -238,7 +238,7 @@ SM 占用率受 registers、shared memory、warp slots 和 CTA slots 的限制 许多现代 Tensor Core kernel 会主动消耗更多资源，即使这会降低 occupancy。多 stage 的 shared memory pipeline 会占用 SMEM；较大的 register fragments 会占用 registers；TMEM allocation 会占用 Tensor Memory 容量；warp specialization 也可能把整组 warp 固定分配给 producer 或 consumer 角色。 -这是有意做出的取舍。这些 kernel 不依靠大量 warp 同时驻留来隐藏延迟，而是在较少的驻留 CTA 内显式重叠不同阶段。只要 pipeline 能让 TMA、Tensor Core 和 store 路径持续运行，低 occupancy 的 kernel 仍然可以获得很高的性能。 +这是有意做出的取舍。这些 kernel 不依靠大量 warp 同时驻留来隐藏延迟，而是在较少的驻留 CTA 内显式重叠不同阶段。只要 pipeline 能让 TMA、Tensor Core 和 store 路径持续运行，低 occupancy 的 kernel 仍可能获得很高的性能。 两种方式各有适用场景。内存访问不规则、难以显式构造流水线的 kernel，通常更依赖高 occupancy；采用深度 staging 和 warp specialization 的 kernel，则可能用较低的 occupancy 换取更充分的阶段重叠。评价一个 kernel 时，不能只看 occupancy，还要看关键硬件单元是否被持续利用。 diff --git a/zh/chapter_tensor_cores/index.md b/zh/chapter_tensor_cores/index.md index 3c5e5185..6ce7670a 100644 --- a/zh/chapter_tensor_cores/index.md +++ b/zh/chapter_tensor_cores/index.md @@ -113,7 +113,7 @@ Blackwell 将长期存活的 accumulator 移入 TMEM。TMEM 是一种作用域 使用 `cta_group::2` 时，MMA 同时访问一个 CTA pair 中两个 CTA 的 TMEM。CTA pair 由同一 cluster 中 `%cluster_ctarank` 仅最低位不同的两个 CTA 组成：其中一个 rank 为偶数，另一个 rank 为奇数。下文分别称它们为偶数 CTA 和奇数 CTA。 -硬件只要求 CTA pair 中的一个 thread 发出 `tcgen05.mma`；指令既可以由偶数 CTA 发出，也可以由奇数 CTA 发出，但 peer CTA 必须仍然处于 active 状态。本书后续的 kernels 通常约定由偶数 CTA 中的一个 thread 发出 MMA，并通过 `tcgen05.commit` 提交完成通知。 +硬件只要求 CTA pair 中的一个 thread 发出 `tcgen05.mma`；指令既可以由偶数 CTA 发出，也可以由奇数 CTA 发出，但 peer CTA 必须仍然处于 active 状态。本书后续的 kernels 通常约定由偶数 CTA 中的一个 thread 发出 MMA，再通过 `tcgen05.commit` 将此前发出的操作的完成状态关联到 `mbarrier`。 Accumulator layout 由 `cta_group`、M 维大小、A 是稠密矩阵还是结构化稀疏矩阵，以及是否使用 `tcgen05.mma.ws` 共同决定。这个 layout 规定逻辑坐标 `(m,n)` 如何映射到 `TLane` 和 `TCol`。 diff --git a/zh/chapter_tirx_layout_api/index.md b/zh/chapter_tirx_layout_api/index.md index 14935140..0c34d9c7 100644 --- a/zh/chapter_tirx_layout_api/index.md +++ b/zh/chapter_tirx_layout_api/index.md @@ -142,7 +142,7 @@ R[2 : 4@warpid] 表示沿 `warpid` 轴放置 2 份副本，两份之间相隔 4 个 warps。 -GPU 硬件经常需要将同一份数据广播到多个 warps、lanes 或存储区域。Replica 直接将这种行为表示为“一项逻辑数据对应多个物理坐标”。 +Replica 表示同一个逻辑元素具有多个物理坐标。它只记录这些副本应当位于哪里；副本如何生成或使用，由实际消费这个 layout 的 tile 操作决定。 ### Offset @@ -406,7 +406,7 @@ TileLayout(S[(8, 64) : (64@m, 1@m)]) m = 64 * i + j ``` -每行包含 64 个 float16，也就是 128 bytes。固定 `j` 并沿 column 读取时，每换一行都会前进 128 bytes，多个访问可能重复落到同一组 banks。 +每行包含 64 个 float16，也就是 128 bytes。如果一组 threads 从不同行读取同一个 column `j`，相邻地址之间会相隔 128 bytes，因此可能反复落到同一组 banks。 Swizzle 让 address 的低位同时依赖较高的 row bits，使原本落到同一个 bank 的 column access 分散到多个 banks。 diff --git a/zh/chapter_tma/index.md b/zh/chapter_tma/index.md index de70ac26..69ecce14 100644 --- a/zh/chapter_tma/index.md +++ b/zh/chapter_tma/index.md @@ -6,7 +6,7 @@ - TMA 负责在 global memory 和 shared memory 之间异步搬运 tile。一个 warp 中只需一个 thread 发起操作，后续的地址计算和数据传输由硬件完成。 - tensor map descriptor 说明 global tensor 如何组织，包括 shape、strides、tile shape 和 swizzle mode；TMA 指令再给出当前 tile 的坐标和 shared-memory 地址。执行 load 时，TMA 可以在写入 shared memory 的同时应用 swizzle，使 tile 直接采用后续 MMA 所需的布局。 -- TMA load 和 store 使用不同的完成通知。Load 通过 `mbarrier` 按已传输的字节数判断数据是否就绪；store 通过 commit group 和 wait group 确认 source buffer 可以复用。 +- TMA load 和 store 使用不同的完成通知。Load 通过 `mbarrier` 登记待传输的字节数，并在该计数归零后确认数据已经就绪；store 通过 commit group 和 wait group 确认 source buffer 可以复用。 ::: 先看 GEMM mainloop 中最常见的场景。Tensor Core 正在计算第 $k$ 个 tile 时，下一组 A、B tiles 必须在当前计算结束前搬入 shared memory。数据如果没有按时到达，Tensor Core 就只能停下来等待，pipeline 中也会出现气泡（pipeline bubble），也就是计算单元因等待数据而空闲的周期。 diff --git a/zh/chapter_tmem/index.md b/zh/chapter_tmem/index.md index 6325d6ac..2fb10607 100644 --- a/zh/chapter_tmem/index.md +++ b/zh/chapter_tmem/index.md @@ -85,7 +85,7 @@ if warp_id == 0: ## 每个 warp 能访问哪些 TMEM lanes -TMEM 属于 CTA，但 `tcgen05.ld` 和 `tcgen05.st` 不会让 CTA 中的任意 warp 访问全部 128 个 Lane 位置。一个 warpgroup 中的四个 warps 各自负责一个由 32 个 Lane 位置组成的固定范围： +TMEM 属于 CTA，但每个 warp 只能通过 `tcgen05.ld` 和 `tcgen05.st` 访问固定的 32-lane window。一个 warpgroup 中的四个 warps 分别覆盖下面四个范围： | warp 在 warpgroup 中的 ID | 可访问的 TMEM Lane 位置 | | --- | --- | diff --git a/zh/index.md b/zh/index.md index 25d94e26..b5a2eb15 100644 --- a/zh/index.md +++ b/zh/index.md @@ -1,14 +1,14 @@ # 面向机器学习系统的现代 GPU 编程 -机器学习系统支撑着现代 AI 的核心计算任务。随着模型规模扩大、部署场景变得更加复杂，系统性能越来越依赖少数关键 GPU kernel 的实现质量。Attention kernel、LLM prefill 和 decode kernel、低精度 block-scaled GEMM、融合 MoE 层，以及其他大型融合 kernel，都会直接影响训练和服务的端到端速度。 +机器学习系统支撑着现代 AI 的许多核心计算任务。随着模型规模扩大、部署场景变得更加复杂，端到端性能越来越依赖少数关键 GPU kernel 的实现质量。Attention、LLM prefill 和 decode、低精度 block-scaled GEMM、融合 MoE 层以及其他大型融合 kernel，都会直接影响训练和服务的速度。 -因此，要理解和优化现代 AI 系统，就必须理解高性能 GPU kernel 是如何写出来的。然而，高性能 kernel 并不是简单堆叠优化技巧的结果。现代 GPU 架构已经发生了显著变化：新的架构引入了更丰富的内存空间、新的数据搬运机制，以及越来越专用化的执行单元。要充分利用这些硬件能力，我们既需要建立清晰的硬件心智模型，也需要理解一个高性能 kernel 是如何从基础版本一步步演化出来的。本书重点关注的正是这两个方面。 +要让这些 kernel 真正跑得快，不能只罗列优化技巧。近年来的 GPU 架构引入了更丰富的内存空间、新的数据搬运机制和越来越专用化的执行单元。要充分利用这些硬件能力，既要理解 GPU 如何执行程序，也要掌握一个基础 kernel 如何逐步演变成高性能实现。本书将围绕这两个方面展开。 -基于这一目标，本书将按照从硬件到代码、再到高性能 kernel 的顺序展开。我们会先介绍 GPU 的硬件组织和执行模型，然后学习本书使用的编程模型，最后在这些基础上逐步构建先进的 GPU kernel。具体来说，本书将以 NVIDIA Blackwell 架构为例，详细讲解 General Matrix-Matrix Multiplication (GEMM)，以及 FlashAttention。在这些 kernel 的构建过程中，我们还会系统学习数据布局、异步数据搬运、异步协作等 GPU 优化中的关键主题。 +本书按照从硬件、编程模型到完整 kernel 的顺序展开。我们会先介绍 GPU 的组织方式和执行模型，再学习本书使用的编程模型，最后逐步构建高性能 kernel。本书主要面向 NVIDIA Blackwell，并以 General Matrix-Matrix Multiplication（GEMM）和 FlashAttention 为贯穿全书的示例。在构建这些 kernel 的过程中，还会系统介绍数据布局、异步数据搬运和异步协作等关键主题。 -本书内容源自卡内基梅隆大学的 [Machine Learning Systems](https://mlsyscourse.org/) 课程系列。为了让这些概念可以通过真实代码学习、运行和验证，本书使用 TIRx Python DSL 逐步构建 GPU kernel 示例。TIRx 贴近硬件，并暴露底层执行抽象，因此读者可以一边运行代码，一边推理其背后的控制流、内存访问和同步逻辑。 +本书内容源自卡内基梅隆大学的 [Machine Learning Systems](https://mlsyscourse.org/) 课程系列。书中的示例使用 TIRx Python DSL，让读者能够在真实 kernel 中学习、运行和验证这些概念。TIRx 会明确表示与硬件执行有关的选择，因此可以结合可运行的代码分析控制流、内存访问和同步逻辑。 -本书是开源项目，欢迎通过 [GitHub 仓库](https://github.com/mlc-ai/modern-gpu-programming-for-mlsys) 提交贡献、修正和示例。 +本书是开源项目，欢迎通过 [GitHub 仓库](https://github.com/mlc-ai/modern-gpu-programming-for-mlsys) 贡献代码、勘误和示例。 ## 本书结构 @@ -19,8 +19,6 @@ - **第四部分：Flash Attention 4。** 这一部分基于第三部分的技术构建完整的 attention kernel：两个 MMA，中间插入 softmax，并包含 online-softmax rescaling、causal mask 和 GQA。 - **参考资料。** TIRx 语言参考、编译器内部机制，以及异步 kernel 调试指南。 -## 已发布章节 - ```{toctree} :caption: 第一部分：理解 GPU :maxdepth: 1 @@ -46,7 +44,7 @@ chapter_tirx_layout_api/index ```{toctree} :caption: 第三部分：GEMM：从 Tiled 到 SOTA -:maxdepth: 1 +:maxdepth: 2 chapter_gemm_basics/index chapter_gemm_async/index @@ -55,7 +53,7 @@ chapter_gemm_advanced/index ```{toctree} :caption: 第四部分：Flash Attention 4 -:maxdepth: 1 +:maxdepth: 2 chapter_flash_attention/index ``` @@ -69,9 +67,3 @@ tirx_guide/language_reference/index appendix/debugging_warp_specialized tirx_guide/arch/index ``` - - diff --git a/zh/tirx_guide/arch/lowering_pipeline.rst b/zh/tirx_guide/arch/lowering_pipeline.rst index 8318b9d8..384258cb 100644 --- a/zh/tirx_guide/arch/lowering_pipeline.rst +++ b/zh/tirx_guide/arch/lowering_pipeline.rst @@ -21,8 +21,8 @@ TIRx Lowering Pipeline 调用 ``tvm.compile(mod, target, tir_pipeline="tirx")`` 时，编译器会将输入的 TIRx module 依次送入一组 TIR passes，这组 passes 称为 **tirx pipeline**。 它负责把 tile primitives、使用 ``TileLayout`` 的 buffers 和 execution-scope -ids 等高层结构，逐步 lowering 为彼此分离的 **host** 与 **device** functions， -最后再由 CUDA backend 生成源码。 +ids 等高层结构逐步转换为彼此分离的 **host** 与 **device** functions，最后再 +由 CUDA backend 生成源码。 Pipeline 定义在 ``python/tvm/tirx/compilation_pipeline.py`` 的 ``tirx_pipeline`` 中。下面按执行顺序介绍其中的 passes。 @@ -55,7 +55,7 @@ Pass 执行顺序 - 作用 * - 1 - ``LowerTIRx`` - - 完成 TIRx 的核心 lowering，详见下方 `LowerTIRx 内部做了什么`_ + - 完成 TIRx 的核心转换，详见下方 `LowerTIRx 内部做了什么`_ * - 2 - ``UnifyThreadBinding`` - 合并等价的 thread-axis bindings，使每个 ``threadIdx`` / ``blockIdx`` @@ -65,7 +65,7 @@ Pass 执行顺序 - 使用 arithmetic analyzer 简化 statement 中的算术表达式 * - 4 - ``LowerTIRxOpaque`` - - 将剩余的 opaque TIRx constructs lowering 为普通 TIR + - 将剩余的 opaque TIRx constructs 转换为普通 TIR * - 5 - ``FlattenBuffer`` - 将多维 ``BufferLoad`` / ``BufferStore`` 展平为一维访问 @@ -113,9 +113,9 @@ Pass 执行顺序 之后，编译器会根据 function 类型分别执行 **finalization**： -- **host**：``LowerTVMBuiltin`` lowering ``tvm_*`` builtins， - ``LowerIntrin`` lowering target-specific intrinsics。 -- **device**：``LowerWarpMemory`` 将 warp-scoped buffers lowering 为 +- **host**：``LowerTVMBuiltin`` 处理 ``tvm_*`` builtins，``LowerIntrin`` + 处理 target-specific intrinsics。 +- **device**：``LowerWarpMemory`` 将 warp-scoped buffers 转换为 shuffles，随后执行 ``StmtSimplify`` 和 ``LowerIntrin``。 LowerTIRx 内部做了什么 @@ -134,7 +134,7 @@ LowerTIRx 内部做了什么 - **``LowerTIRxCleanup``** 运行 ``LayoutApplier``，将使用 ``TileLayout`` 的 buffer access 变成具体的物理地址计算 （``addr = data + elem_offset + layout.apply(coord)``），再展平 buffers， - 并将 execution-scope ids lowering 为 thread axes，例如 + 并将 execution-scope ids 转换为 thread axes，例如 ``T.cta_id`` / ``T.thread_id`` 通过 ``launch_thread`` 变为 ``blockIdx`` / ``threadIdx``。 diff --git a/zh/tirx_guide/language_reference/cuda/buffers.rst b/zh/tirx_guide/language_reference/cuda/buffers.rst index 695d9676..13f4e435 100644 --- a/zh/tirx_guide/language_reference/cuda/buffers.rst +++ b/zh/tirx_guide/language_reference/cuda/buffers.rst @@ -210,10 +210,10 @@ dynamic shared memory allocation。Static shared memory 的大小在编译期 .. code-block:: python - # device kernel attribute: + # device kernel 属性： "tirx.kernel_launch_params": ["blockIdx.x", "threadIdx.x", "tirx.use_dyn_shared_memory"] - # host-side launch call (..., gridDim.x, blockDim.x, dyn_shared_bytes): + # host 侧 launch 调用（..., gridDim.x, blockDim.x, dyn_shared_bytes）： T.call_packed("dyn_kernel", A.data, B.data, C.data, 1, 64, 512) 运行时，这里的 ``512`` 会成为 ``cuLaunchKernelEx`` 调用中的 @@ -230,28 +230,29 @@ helper，以及将 cursor 回退以复用空间的 ``move_base_to``： .. code-block:: python - pool = T.SMEMPool() # bump allocator over shared.dyn - As = pool.alloc((BM, BK), "float16", align=128) # carve a tile + pool = T.SMEMPool() # shared.dyn 上的 bump allocator + As = pool.alloc((BM, BK), "float16", align=128) # 分配一个 tile Bs = pool.alloc((BK, BN), "float16", align=128) - Cs = pool.alloc_mma((BM, BN), "float16") # MMA-compatible, swizzle inferred - pool.commit() # finalize the pool's size - # pool.move_base_to(offset) rewinds the cursor to reuse space + Cs = pool.alloc_mma((BM, BN), "float16") # 自动推导 MMA-compatible swizzle + pool.commit() # 确定 pool 的最终大小 + # pool.move_base_to(offset) 将 cursor 回退到可复用的位置 下方的 TMEM pool 建立在 ``SMEMPool`` 之上。 Registers --------- -Per-thread 临时数据位于 registers 中。使用 -``T.alloc_local(shape, dtype)``（即 ``scope="local"``）分配；它只属于当前 -thread，并会生成保存在 registers 中的 local array。 +Per-thread 临时数据使用 ``local`` scope。通过 +``T.alloc_local(shape, dtype)`` 分配后，这些数据只属于当前 thread。使用 +静态索引的 local arrays 通常会被 scalarize 到 registers；使用动态索引，或 +register pressure 较高时，也可能进入 local memory。 .. code-block:: python - r = T.alloc_local((4,), "float32") # per-thread register array + r = T.alloc_local((4,), "float32") # 每个 thread 私有的 register array for k in T.unroll(4): r[k] = A[tx, k] - # ... compute on r[0..3] ... + # ... 使用 r[0..3] 计算 ... .. code-block:: c++ @@ -281,7 +282,7 @@ Scalar 本质上是只有**一个元素**的 register array。可以直接分配 .. code-block:: python - phase = T.alloc_local((1,), "int32") # 1-element register array + phase = T.alloc_local((1,), "int32") # 单元素 register array phase[0] = 0 while phase[0] < 4: acc = acc + A[tx, phase[0]] @@ -292,20 +293,20 @@ Scalar 本质上是只有**一个元素**的 register array。可以直接分配 .. code-block:: python - phase: T.int32 = 0 # mutable scalar (sugar for the above) + phase: T.int32 = 0 # mutable scalar，是上一种写法的语法糖 while phase < 4: acc = acc + A[tx, phase] phase += 1 - s = T.local_scalar("int32") # explicit form; assign by name (s = ..., not s[0]) - acc: T.float32 = 0.0 # a type-annotated assignment also makes one + s = T.local_scalar("int32") # 显式形式；通过名称赋值，而不是 s[0] + acc: T.float32 = 0.0 # 带类型注解的赋值也会创建 scalar 两种写法在 parse 后会得到结构完全相同的 TIRx。Parser 会将 ``phase: T.int32`` 解析为单元素 ``local`` buffer，将 ``phase`` 和 ``phase += 1`` 解析为 ``phase[0]`` 和 ``phase[0] += 1``。对两个 kernels 调用 ``tvm.ir.assert_structural_equal`` 会通过；printer 甚至会把显式的 ``alloc_local`` 加 ``[0]`` 重新输出为 scalar 语法。因此，parse 完成后两者 -没有区别，都会 lowering 为 ``alignas(64) int phase_ptr[1];``。Scalar +没有区别，都会生成 ``alignas(64) int phase_ptr[1];``。Scalar 只是省去了 ``[0]``。``T.local_scalar``、``T.shared_scalar`` 和 ``T.alloc_scalar`` 可以显式选择 scope。 @@ -324,15 +325,15 @@ Scalar 本质上是只有**一个元素**的 register array。可以直接分配 .. code-block:: python - n: T.let = M * K # immutable binding (LetStmt) - half: T.let[T.int32] = N // 2 # ... with an explicit type + n: T.let = M * K # immutable binding（LetStmt） + half: T.let[T.int32] = N // 2 # 显式指定类型 它会生成普通的 C scalar variable，而不是 array，也不需要 ``[0]``。例如， 运行时变量 ``m`` 上的 ``half: T.let = m * 2`` 会生成： .. code-block:: c++ - int half = m * 2; // the `let` -> a const-like local + int half = m * 2; // `let` 生成类似 const 的 local variable 由于值不会改变，simplifier 可以自由执行 propagation 和 common subexpression elimination。因此在使用位置可能直接看到 ``m * 2``，也可能 @@ -368,15 +369,15 @@ offset 上 ``decl`` tensor views，结束时由一个 warp 释放： .. code-block:: python - addr = T.alloc_shared((1,), "uint32") # slot for the allocated base - if warp_id == alloc_warp: # tcgen05.alloc is warp-uniform + addr = T.alloc_shared((1,), "uint32") # 保存 allocation base 的 slot + if warp_id == alloc_warp: # tcgen05.alloc 是 warp-uniform T.ptx.tcgen05.alloc(T.address_of(addr), n_cols=512, cta_group=cta_group) acc = T.decl_buffer((CTA_M, 512), "float32", scope="tmem", - allocated_addr=0, layout=tmem_layout) # view at column 0 - # ... use acc as a gemm_async / copy_async operand ... + allocated_addr=0, layout=tmem_layout) # column 0 处的 view + # ... 将 acc 用作 gemm_async / copy_async operand ... if warp_id == alloc_warp: T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) - T.ptx.tcgen05.dealloc(addr, n_cols=512, cta_group=cta_group) + T.ptx.tcgen05.dealloc(addr[0], n_cols=512, cta_group=cta_group) 此时 column offsets 和 ``tmem_layout`` （datapath D/F layout）都需要手工 管理。下面的 pool 会自动生成同样的步骤。 @@ -389,13 +390,13 @@ datapath layout： .. code-block:: python - tmem_addr = pool.alloc((1,), "uint32") # pool = the kernel's smem pool + tmem_addr = pool.alloc((1,), "uint32") # pool 是 kernel 的 SMEM pool tmem_pool = T.TMEMPool(pool, total_cols=512, cta_group=cta_group, tmem_addr=tmem_addr) - acc = tmem_pool.alloc((CTA_M, 512), "float32") # allocated_addr set for you - tmem_pool.commit() # emits tcgen05.alloc (one warp) - # ... use acc ... - tmem_pool.dealloc() # emits tcgen05.dealloc (one warp) + acc = tmem_pool.alloc((CTA_M, 512), "float32") # 自动设置 allocated_addr + tmem_pool.commit() # 由一个 warp 发出 tcgen05.alloc + # ... 使用 acc ... + tmem_pool.dealloc() # 由一个 warp 发出 tcgen05.dealloc 完整示例见第三部分的 GEMM kernels。 diff --git a/zh/tirx_guide/language_reference/cuda/control_flow.rst b/zh/tirx_guide/language_reference/cuda/control_flow.rst index 65bbebb3..028f124c 100644 --- a/zh/tirx_guide/language_reference/cuda/control_flow.rst +++ b/zh/tirx_guide/language_reference/cuda/control_flow.rst @@ -18,8 +18,8 @@ 控制流 ====== -TIRx 的控制流包括 ``if``、多种 loop 和 ``while``，它们会生成对应的 CUDA -控制流。 +TIRx 提供 ``if``、多种 loop 和 ``while``，它们会直接映射到对应的 CUDA +控制流结构。 if -- @@ -46,9 +46,9 @@ warp 中选出一个 thread 发出指令： A_ptr[tx] = A_ptr[tx] + 1.0f; } -如果只需要在表达式中选择值，不需要控制流分支，可以使用 -``T.if_then_else(cond, a, b)``。它会 lowering 为三元表达式，不会产生 -control-flow divergence： +如果需要在表达式中选择值，而不在 TIRx 中建立显式控制流分支，可以使用 +``T.if_then_else(cond, a, b)``。它会转换为三元表达式；最终使用哪些机器 +指令实现这个表达式，仍由 backend 决定： .. code-block:: c++ @@ -67,10 +67,11 @@ Uniform 与 Divergent 控制流 ``T.cuda.warpgroup_sync(id)``，详见第三部分的 warp-specialized GEMM 和 :doc:`threads_sync`。 -初始化 barrier 时也要注意参与范围。``mbarrier`` 的 ``.init()`` 会 lowering -为 single-thread guard（``if (threadIdx.x < 1)``）。如果再把它放进另一个 -divergent branch，barrier 可能没有被初始化，进而导致 unspecified launch -failure。 +初始化 barrier 时也要注意参与范围。高层 ``MBarrier.init()`` wrapper 会生成 +single-thread guard（``if (threadIdx.x < 1)``）。如果再把它放进另一个 +divergent branch，barrier 可能没有初始化，进而导致 unspecified launch +failure。原始的 ``T.ptx.mbarrier.init`` intrinsic 不会自动添加这个 guard； +调用者必须自行选出负责初始化的 thread。 loop ---- @@ -110,7 +111,7 @@ while A[i] = A[i] + T.float32(1.0) i += 1 -它会 lowering 为带有提前退出 ``break`` 的 ``while (1)``。其中计数器使用 +它会转换为带有提前退出 ``break`` 的 ``while (1)``。其中计数器使用 一个只有一个元素的 register buffer： .. code-block:: c++ diff --git a/zh/tirx_guide/language_reference/cuda/threads_sync.rst b/zh/tirx_guide/language_reference/cuda/threads_sync.rst index 662a2be3..e66c1ee1 100644 --- a/zh/tirx_guide/language_reference/cuda/threads_sync.rst +++ b/zh/tirx_guide/language_reference/cuda/threads_sync.rst @@ -56,7 +56,7 @@ CUDA C++/PTX Intrinsics i[0] = i[0] // 2 A[lane_id] = v[0] -Shuffle 会直接 lowering 为 ``__shfl_xor_sync``： +Shuffle 会直接转换为 ``__shfl_xor_sync``： .. code-block:: c++ @@ -64,22 +64,22 @@ Shuffle 会直接 lowering 为 ``__shfl_xor_sync``： ``T.ptx.*`` / ``T.cuda.*`` 还包含 ``cp_async``（LDGSTS）、 ``cp_async.bulk.tensor``（TMA）、``ldmatrix`` / ``stmatrix``、``tcgen05.*`` -（Blackwell MMA）、``atomic_add`` 和 ``fence`` 等 families。完整列表请参阅 +（Blackwell MMA）、``atomic_add`` 和 ``fence`` 等指令类别。完整列表请参阅 ``tvm.backend.cuda`` backend API reference。 同步语义 -------- GEMM 和 Flash Attention kernel 中经常出现下面四种同步机制。它们控制异步 -engine 和并行 thread groups，使用错误通常会导致 silent corruption 或 -deadlock。 +engine 和并行 thread groups；使用错误通常会让结果在没有报错的情况下损坏， +或者造成 deadlock。 **Mbarrier phase。** Mbarrier 使用一个内部 phase bit 追踪不同轮次的 arrival。``T.ptx.mbarrier.try_wait(bar, phase)`` 会一直等待，直到 barrier 内部 phase 与调用者提供的 ``phase`` 不同。循环复用 barrier 时，每次 wait -之后都必须翻转本地 phase tracker（``phase ^= 1``）。否则，后续 wait 会立即 -返回，engine 可能读取只写完一部分的 memory。第三部分的 GEMM 章节会给出 -完整的 phase tracking 表。 +之后都必须翻转本地 phase tracker（``phase ^= 1``）。否则，后续 wait 可能 +误把上一轮的完成状态当作当前轮，导致 consumer 在本轮 producer 或异步操作 +真正完成前访问数据。第三部分的 GEMM 章节会给出完整的 phase tracking 表。 **Election。** ``T.ptx.elect_sync()`` 从一个 warp 的 active lanes 中选择 **一个 lane**；它不一定选择 lane 0，也不是每个 CTA 选择一个 thread。如果 @@ -92,8 +92,8 @@ kernel 使用 ``if warp_id == 0:`` 加 ``if T.ptx.elect_sync():`` 发出 角色分支后，不能把 ``cta_sync()`` 放进其中一个 warpgroup 的分支，否则其他 warpgroups 无法到达，kernel 会 deadlock。硬件提供 16 个 named barriers （ID 0 到 15）；``T.cuda.warpgroup_sync(10)`` 只同步一个 warpgroup 的 -threads。不同 warpgroups 使用不同 ID，例如 -``warpgroup_sync(wg_id + 10)``，避免共享同一个 hardware barrier。第三部分 +threads。可能同时处于 active 状态的独立同步必须使用不同 ID，例如 +``warpgroup_sync(wg_id + 10)``；前一次同步完成后，ID 可以再次使用。第三部分 的 warp-specialized GEMM 会展示完整用法。 **Fence。** Fence 保证 producer 的写入先于 consumer（通常是异步 engine）