Polish and align English and Chinese documentation

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 |

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`).

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.

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,

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.