B200 Tmem Lifecycle Planner

B200 TMEM Lifecycle Planner

R — Source evidence (Reading, paraphrased)

• [S8] TMEM is a Blackwell-only, CTA-scoped two-dimensional Lane × Col scratchpad, with 128 Lane rows and up to 512 32-bit Cols.
• [S8] TMEM must be explicitly allocated and freed in 32-column units; ordinary shared load/store cannot access it.
• [S8] tcgen05.ld/st/cp are all dedicated asynchronous paths, each with its own completion rules.
• [S7] Both the accumulator and the block-scale factors may live in TMEM, but their layouts and lifecycles differ.

Source: distilled from "Modern GPU Programming for MLSys" (https://mlc.ai/modern-gpu-programming-for-mlsys/) and the NVIDIA Blackwell tuning/compatibility guides. Short paraphrases only; no long passages are reproduced.

I — Methodology skeleton (Interpretation)

TMEM planning resembles an "explicitly managed accumulator heap", not compiler-allocated registers:

1. First compute the columns and alignment each object class needs;
2. then map the logical tiles onto TLane/TCol;
3. annotate who writes, who reads, and when each completes;
4. arrange region reuse and the final free;
5. ensure the epilogue's register fragments are consistent with the TMEM mapping used at write time.

A larger tile reducing register pressure does not mean the resources are free; it shifts the pressure onto the TMEM column budget and readback bandwidth.

A1 — Applications in the source (Past Application)

Case 1: dense MMA accumulator

• tcgen05.mma writes the fp32 accumulator to TMEM.
• After the compute phase, the full warpgroup uses tcgen05.ld to read its fragments back into registers for cast/epilogue/store.

Case 2: block-scaled MMA

• A/B data live in SMEM.
• Scale factors first go to SMEM, then enter the dedicated TMEM scale layout via tcgen05.cp.
• The scale region and the accumulator region must not reuse the same mapping just because both are in TMEM.

A2 — Trigger scenarios (Future Trigger) ★

In what situations will the user need this skill?

1. "How many TMEM columns does this tcgen05 kernel need to allocate?"
2. "Help me plan the TMEM reuse of S/P/O in FlashAttention."
3. "How should the completion and free ordering of tcgen05.ld/st/cp be written?"

Language signals

• "How many TMEM columns does this tcgen05 kernel need to allocate?"
• "Help me plan the TMEM reuse of S/P/O in FlashAttention."
• "How should the completion and free ordering of tcgen05.ld/st/cp be written?"

Distinction from adjacent skills

Versus b200-tcgen05-mma-contract-builder: this skill focuses on TMEM resources, layout, and lifecycle; the MMA contract builder additionally decides the tile shape, cta_group, dtype, and operand placement.

E — Executable steps (Execution)

Once the skill is activated, the agent must execute the following procedure:

1. List the TMEM objects
   • Accumulators, score/prob/output tiles, scale factors, temporary regions.
   • For each object, record its dtype, logical shape, and the number simultaneously live.
2. Compute the column budget
   • Map the objects onto 32-bit Cols; round up to the 32-column allocation unit.
   • Include the replication factors for pipeline stages, dual Q stages, or multiple consumers.
3. Define the 2D layout
   • Give explicit TLane, TCol formulas; "so many bytes off the base" alone is not enough.
   • For cta_group::2, write out each CTA's fragment in its own TMEM separately.
4. Define the write paths
   • MMA accumulator: tcgen05.mma.
   • Scale: tcgen05.cp from SMEM→TMEM.
   • register→TMEM: tcgen05.st.
5. Define the readback path
   • Choose the tcgen05.ld shape/repeat; write out the fragment each warp/lane receives.
   • After readback, wait on the corresponding wait::ld before consuming the registers.
6. Draw the lifecycle
   • allocate → producer writes → completion → consumers read/modify → completion → reuse/free.
   • Annotate every reuse edge with its barrier or wait.
7. Check region reuse
   • Overlay only when all old consumers have completed and the new layout does not conflict.
8. Produce the resource report
   • Total columns, peak simultaneously-live columns, alignment waste, the riskiest reuse edge, and the epilogue readback cost.

Required outputs

1. Conclusion: the current choice/diagnosis, without vague "could be any of them" hedging.
2. Evidence or assumptions: which come from user data, and which are hypotheses awaiting verification.
3. Contract/table/timeline: the auditable intermediate artifacts corresponding to this skill.
4. Minimal validation: correctness tests, boundary tests, and one falsifiable experiment.
5. Risks and fallback: alternative paths when hardware, version, or resource conditions are not met.

B — Boundaries (Boundary) ★

Do not use when

• The target is not Blackwell, or the kernel does not use tcgen05/TMEM.
• The user only needs ordinary shared memory tiling.

Failure modes

• Accessing TMEM with ld.shared/st.shared.
• Forgetting to allocate in 32-column units, or forgetting to free.
• Consuming registers before tcgen05.ld has completed.
• Mixing the accumulator layout with the scale-factor layout.

Limitations

• The exact available column count, supported load shapes, and lowering depend on the toolchain version; verify against the target compiler and the generated PTX.

Related skills

• depends-on: b200-layout-contract-auditor
• contrasts-with: none
• composes-with: b200-tcgen05-mma-contract-builder, b200-mbarrier-protocol-auditor, b200-flash-attention4-planner

Audit info

• Validation passed: V1 ✓ / V2 ✓ / V3 ✓
• Test definitions: 6 (3 should_trigger / 2 should_not_trigger / 1 edge_case)
• Hardware validation: not performed; must be verified on a target B200
• Distilled: 2026-06-25