# TMEM Round-Trip Investigation Plan

**Purpose:** Determine whether ROADMAP Priority 8, Path A (CUTLASS atom replication) is achievable in CuTeDSL Python. If yes, this unlocks single-launch multi-tile FMHA. If no, we fall back to Path C (O in SMEM) with documented evidence.

**Premise:** CUTLASS C++ Blackwell FMHA performs lossless TMEM round-trip for O rescale using `SM100_TMEM_LOAD_32dp32b_*x` and `SM100_TMEM_STORE_32dp32b_*x` paired atoms. The hardware supports this. The question is whether CuTeDSL Python exposes the matching atom + layout configuration.

**Time budget:** 5 working days. If NO-OP round-trip can't be made bitwise-correct in that window, escalate to Path C.

---

## What we know vs what we don't

### Known facts

- The hardware supports lossless TMEM round-trip — CUTLASS C++ FMHA does it in production for correction_rescale.
- Current FMHA code at `dsv4/kernels/attention/fmha.py` lines 407–416 uses `tcgen05.copy.Ld32x32bOp(Repetition(16))` + `tcgen05.copy.St32x32bOp(Repetition(16))` built independently from the same `tOtO_i` tensor via `make_tmem_copy`.
- That configuration produces ~3% error on NO-OP round-trip (`test_d1_tmem_trip.py`).
- The MoE kernel never does a TMEM round-trip, so the MoE epilogue pattern is not a template here.
- The S load uses `Ld32x32bOp(Repetition(32))` and works correctly for the QK→softmax path (no round-trip required — load only).

### Open questions (in priority order)

1. **Does CUTLASS publish a CuTeDSL Python FMHA reference that does in-kernel O rescale?** If yes, the answer is in that file. (Path noted in README: `/root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py`.)
2. What specific atom variants does CUTLASS C++ use for the round-trip — `32dp32b_1x` / `4x` / `16x` / etc., transposed or non-transposed?
3. Does the FP32 accumulator layout require a specific Repetition value to round-trip cleanly?
4. Are the load/store atoms supposed to be constructed from the same partitioning (`make_tmem_copy(load_atom, T)` and `make_tmem_copy(store_atom, T)` with the same T), or from each other via some derivation we haven't found?
5. Is there an intermediate layout transformation (analogous to `transform_partitioned_tensor_layout`) that's required before round-tripping?
6. Does CuTeDSL Python expose all the atom variants CUTLASS C++ uses? If not, which are missing?

---

## Phase 1: Read the CUTLASS reference (Day 1, half day)

**Goal:** Find the exact rescale code in CUTLASS and document its atom + partition + copy pattern.

### Steps

- [ ] Locate `/root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py` on B200. Confirm it exists.
- [ ] Search for the correction/rescale code path. Likely keywords: `correction`, `rescale`, `acc_scale`, `row_max`, `softmax_correction`.
- [ ] Extract the rescale block end-to-end. Copy it into `notes/cutlass_rescale_excerpt.py` (gitignored).
- [ ] Document each variable in that block:
  - Source tensor for `make_tmem_copy`
  - Atom name + parameters
  - Repetition factor
  - Thread/partition layout
  - Any layout transforms applied before copy
  - Fence calls
- [ ] **Then look at the C++ reference for the same code path.** Local clone at `/home/openclaw/dev/cutlass`. File: likely under `include/cutlass/gemm/collective/sm100_fmha_*` or `examples/77_blackwell_fmha/`. The C++ may be more readable than the Python in places — the Python is a port of the C++ semantics.

### If CuTeDSL Python reference doesn't exist or doesn't do rescale

This is the "early bad news" outcome. Document it and escalate immediately — the entire premise of Path A becomes shakier. Still worth Phase 2, but treat it as exploratory rather than porting.

### Deliverable

`notes/cutlass_rescale_breakdown.md` — a side-by-side: CUTLASS pattern vs current FMHA pattern, with the specific differences highlighted. This file is gitignored (private notes) but its contents drive Phase 2.

---

## Phase 2: Standalone NO-OP round-trip test (Days 1–2)

**Goal:** Build the minimal kernel that does *only* TMEM round-trip on a known FP32 buffer, and get cos = 1.0 (bitwise identity).

This is the gating test. Until NO-OP round-trip passes in isolation, no amount of FMHA-specific rescale logic helps.

### Minimum viable test kernel

Create `tests/unit/test_tmem_roundtrip_minimal.py` with:

```python
# Pseudocode — exact API verified on B200
class TmemRoundtripKernel:
    def __init__(self, atom_load, atom_store, n_rows=128, n_cols=64):
        self.atom_load = atom_load
        self.atom_store = atom_store
        self.n_rows = n_rows
        self.n_cols = n_cols

    @cute.jit
    def __call__(self, src, dst, stream):
        # Launch single CTA with 4 warps (matching FMHA epilogue warp count)
        self._kernel(src, dst).launch(grid=(1,1,1), block=[128,1,1], stream=stream)

    @cute.kernel
    def _kernel(self, src, dst):
        # 1. Allocate TMEM
        tmem.allocate(power_of_2_round_up(self.n_cols))
        tmem_ptr = tmem.retrieve_ptr(Float32)

        # 2. Build a TMEM tensor at known offset
        layout = make_tmem_layout(self.n_rows, self.n_cols)
        t_tensor = cute.make_tensor(tmem_ptr, layout)

        # 3. Write known values to TMEM via store atom
        #    (need a way to seed TMEM — could come from MMA, or from a manual SMEM->TMEM copy)

        # 4. Load TMEM -> registers via atom_load
        load_partition = make_tmem_copy(self.atom_load, t_tensor)
        thr_load = load_partition.get_slice(tidx)
        reg_buf = cute.make_rmem_tensor(thr_load.partition_D(...).shape, Float32)
        cute.copy(load_partition, thr_load.partition_S(t_tensor), reg_buf)
        cute.arch.fence_view_async_tmem_load()

        # 5. NO-OP — don't modify reg_buf

        # 6. Store registers -> TMEM via atom_store
        store_partition = make_tmem_copy(self.atom_store, t_tensor)
        thr_store = store_partition.get_slice(tidx)
        cute.copy(store_partition, reg_buf, thr_store.partition_D(t_tensor))
        cute.arch.fence_view_async_tmem_store()

        # 7. Load TMEM -> registers again via same load atom
        # 8. Write registers -> GMEM dst for comparison
```

The challenge in step 3 is **seeding TMEM with known values**. Options:
- **A:** Run a small MMA with known inputs to populate the TMEM accumulator. Most realistic (matches FMHA usage).
- **B:** Manual SMEM → TMEM copy via `tcgen05.cp` with known SMEM contents. More direct.
- **C:** Use a separate "init" kernel that populates TMEM, then the round-trip kernel reads/writes it. Awkward but possible.

Recommend **B** for the minimal test — gives full control over the seed values, no MMA semantics in the way.

### Atom variant matrix to try

Start with what CUTLASS uses (extracted in Phase 1) and the current FMHA configuration. Iterate from there.

| Variant | Atom (Load) | Atom (Store) | Repetition | Source Tensor | Notes |
|---|---|---|---|---|---|
| V1 (baseline) | `Ld32x32bOp` | `St32x32bOp` | 16 | `tOtO_i` (composition-tiled) | Current FMHA — known to fail |
| V2 (CUTLASS-style) | from Phase 1 | from Phase 1 | from Phase 1 | from Phase 1 | The target |
| V3 (Repetition mismatch?) | `Ld32x32bOp` | `St32x32bOp` | 32 / 32 | `tOtO0` (no composition) | Less tiled |
| V4 (different N) | `Ld16x32bOp` if exists | `St16x32bOp` if exists | varies | as above | Other datapath count |
| V5 (Nt variant) | `Ld32x32bNtOp` if exists | `St32x32bNtOp` if exists | from Phase 1 | as above | Transposed variant |
| V6 (no composition) | as V1 | as V1 | 16 | raw `tOtO0` | Test if `composition` is the culprit |

### Test protocol per variant

For each variant:

1. Print the partition shapes at trace time for both load and store:
   ```python
   print(f"V{N}: load partition_S shape = {cute.shape(thr_load.partition_S(t_tensor))}")
   print(f"V{N}: load partition_D shape = {cute.shape(thr_load.partition_D(coord_tensor))}")
   print(f"V{N}: store partition_S shape = {cute.shape(thr_store.partition_S(coord_tensor))}")
   print(f"V{N}: store partition_D shape = {cute.shape(thr_store.partition_D(t_tensor))}")
   ```
2. Seed TMEM with `arange(n_rows * n_cols)` reshaped to `(n_rows, n_cols)` — predictable, easy to spot specific corruption patterns.
3. Run NO-OP round-trip.
4. Read TMEM back, compare bitwise (`torch.equal`, not cosine) against the seeded values.
5. If not bitwise equal, compute element-by-element diff. Look for patterns:
   - Are wrong values consistent (suggests a layout mismatch — e.g., rows or columns are permuted)?
   - Are wrong values random-looking (suggests races or addressing errors)?
   - Are some columns/rows correct and others wrong (suggests partial coverage)?
6. Record the result in a table `notes/tmem_roundtrip_matrix.md`.

### Decision point

- **NO-OP passes bitwise on any variant:** ✅ Phase 2 successful. Move to Phase 3 with that variant.
- **NO-OP fails on all tested variants, but ≥1 produces a recognizable pattern** (e.g., transposed values, rotated columns): there's likely a missing layout transform. Add a layout transform step and try again.
- **NO-OP fails on all variants with random-looking corruption:** the atom is not actually paired with the variants CuTeDSL Python exposes. Escalate.

---

## Phase 3: Diagnose layout differences (Day 3)

**Only enter Phase 3 if Phase 2 had a pattern but not bitwise success.**

**Goal:** Identify the specific layout transformation needed to make the round-trip lossless.

### Investigation steps

- [ ] Print the TMEM thread-data mapping for both load and store atoms:
  ```python
  print(f"load thr_idx={sfw_idx} → cols {thr_load.partition_S(t).layout}")
  print(f"store thr_idx={sfw_idx} → cols {thr_store.partition_D(t).layout}")
  ```
- [ ] If they differ, find the transformation that maps load layout to store layout.
- [ ] Try inserting a register-level permutation between load and store that compensates for the column mapping difference.
- [ ] Compare against CUTLASS reference: are they doing such a permutation? Is it implicit in some helper?
- [ ] Check whether `transform_partitioned_tensor_layout` (used in MoE epilogue for one-way) has a round-trip-friendly counterpart.
- [ ] Examine `tcgen05.copy.Repetition`'s semantics — does the repetition factor affect column striding in a way that load/store interpret differently?

### Specific things to try

1. **Inverse-permute between load and store.** If load reads `(thread t, col c) → tmem[t, σ(c)]` and store writes `(thread t, col c) → tmem[t, τ(c)]`, with σ ≠ τ, then inserting a register permutation `reg[c'] = reg[σ⁻¹(τ(c'))]` between load and store should fix it.
2. **Build the store TiledCopy from the load's partition object** rather than from the tensor directly. If CuTeDSL exposes a way to derive a sister TiledCopy from an existing one (`.flip()`, `.invert()`, parameter swap), this is where it would live. (My earlier suggestion of `retile_to_S()` was fabricated — check the actual API.)
3. **Apply `cute.composition` or `cute.logical_divide` to the source tensor** before constructing the store atom. README note 10 flags that these produce different layouts.

### Deliverable

If Phase 3 yields bitwise NO-OP success: `notes/tmem_roundtrip_solution.md` documenting the exact atom + partition + transform + permutation combination. This becomes the FMHA fix.

---

## Phase 4: Apply to FMHA (Days 4–5)

**Only enter Phase 4 if Phase 2 or 3 produced bitwise NO-OP success.**

### Steps

- [ ] Replace the rescale atoms in `dsv4/kernels/attention/fmha.py` lines 407–416 with the verified Phase 2/3 configuration.
- [ ] Run hd=64 single-tile regression. Verify cos ≥ 0.999998.
- [ ] Run hd=128 single-tile regression. Verify cos ≥ 0.999997.
- [ ] Run hd=256 single-tile regression. Verify cos ≥ 0.999998.
- [ ] Add a new test: `tests/unit/test_d15_in_kernel_rescale.py` that runs s_k=256, 384, 512, 1024 in a single kernel launch with in-kernel rescale, and compares against Python KV merge oracle.
- [ ] If new test passes at all s_k values: ✅ Path A complete. Remove `n_kv_tiles > 1` guards or change their semantics to "use in-kernel rescale" instead of "skip rescale entirely."
- [ ] If new test fails at higher s_k but passes at s_k=256: there's an interaction between rescale and the pipeline state cycling. Diagnose.

### Test sanity

Bitwise NO-OP in a minimal kernel does not automatically mean correctness inside FMHA. The FMHA has additional concurrent activity (TMA warp, MMA warp, softmax warps all touching TMEM regions). Watch for:
- Fence ordering: every TMEM store must be followed by `fence_view_async_tmem_store()`. Every load must be preceded by `fence_view_async_tmem_load()` if a recent store could affect it.
- Race conditions between the rescale and the next MMA write to the same TMEM region.
- Pipeline state: the rescale happens inside the kt loop after softmax_done_bar but before the next QK MMA. Verify the barrier ordering keeps rescale and MMA from overlapping on O.

---

## Phase 5 (optional): Document and contribute upstream (Day 5+)

If Phase 4 succeeds, the solution is potentially useful to other CuTeDSL Python users.

- [ ] Write `notes/tmem_roundtrip_findings.md` with: the broken pattern, the working pattern, why they differ, what CUTLASS does that we ported.
- [ ] Consider opening an issue / discussion on the CUTLASS repo if the working pattern requires non-obvious atom construction. The CuTeDSL Python documentation may benefit from a worked example.
- [ ] Cross-reference from README "Lessons learned" section.

---

## Escalation triggers (when to stop and switch paths)

The investigation has hard exit criteria so we don't sink unbounded time into it.

| Trigger | Action |
|---|---|
| End of Day 1: CUTLASS CuTeDSL Python FMHA reference doesn't exist or doesn't do rescale | Continue to Phase 2 as exploratory, expect lower success probability. Plan check-in at end of Day 3. |
| End of Day 2: no variant produces NO-OP success, no clear pattern in failures | Stop. Path A is likely not viable in current CuTeDSL Python. Escalate to Path C (O in SMEM). |
| End of Day 3: Phase 3 diagnostic shows the required atom or transform doesn't exist in CuTeDSL Python | Stop. Document the gap. File issue with NVIDIA. Escalate to Path C. |
| End of Day 4: NO-OP passes in isolation but FMHA regression fails and root cause isn't fence ordering | Stop. Path A is partially viable but FMHA-specific interactions are blocking. Document and escalate to Path C while keeping the Phase 2/3 atom config as a future re-attempt seed. |
| End of Day 5: all FMHA regressions pass, new multi-tile test passes | ✅ Path A complete. Proceed to ROADMAP Priority 8 cleanup (delete Python KV merge code, update tests). |

---

## What to leave behind

Whether Path A succeeds or fails, the investigation should produce:

- [ ] `tests/unit/test_tmem_roundtrip_minimal.py` — the standalone NO-OP test, parametrized over atom variants. Future work can re-run this if CuTeDSL Python adds new atoms.
- [ ] `notes/tmem_roundtrip_matrix.md` — the variant test results table.
- [ ] `notes/cutlass_rescale_breakdown.md` — the CUTLASS reference excerpt and analysis.
- [ ] If solved: updated `dsv4/kernels/attention/fmha.py` + new `tests/unit/test_d15_in_kernel_rescale.py`.
- [ ] If unsolved: documented evidence in ROADMAP that Path A was investigated and failed, with the specific gap identified, so future revisitations don't redo the work.

---

## What this investigation deliberately does NOT do

- It does not refactor FMHA's PV accumulator (that's Path B/C, separate work).
- It does not touch the per-head Python launch or multi-CTA grid (Priority 4, depends on Priority 2 not on this).
- It does not change the un-normalized O + LSE output contract (that's exact and ships today).
- It does not address the hd=512 MLIR hang (Priority 9, unrelated toolchain issue).

If during the investigation we find ourselves modifying any of these, stop and reconsider scope. The investigation's win condition is narrow: NO-OP TMEM round-trip works, and FMHA's existing rescale code path can use it without other changes.