nvfp4-megamoe-kernel/STAGE_D1.3.md

STAGE_D1.3.md — SMEM-P Manual Addressing Notes

Problem Statement

Copy P (attention probabilities) from TMEM to SMEM with layout conversion:

• Source: TMEM layout from QK C-fragment
• Destination: SMEM layout for PV A-operand
• Both represent same 128×128 P matrix, different tiling

Layouts Observed (from debug prints)

QK C-fragment Layout (TMEM)

tStS0 layout: ((128,128),1,1):((65536,1),0,0)
tStS0 shape: ((128, 128), 1, 1)

Interpretation: Shape (128,128) with stride (65536, 1) — column-major? Stride between rows = 65536, between columns = 1.

PV A-operand SMEM Layout

sP layout: ((128,16),1,(4,2),1):((64,1),0,(16,8192),0)
sP shape: ((128, 16), 1, (4, 2), 1)

Interpretation: Complex tiling with modes: (128,16), 1, (4,2), 1

Failed Approaches

1. make_tiled_copy_C (Helpers are a trap)

• Created tiled_smem_copy = cute.make_tiled_copy_C(smem_copy_atom, qk_mma)
• partition_S and partition_D produce incompatible tensors
• Source partition: 65536 elements, rank 4
• Destination partition: 2097152 elements, rank 5 (32× mismatch)
• Confirmed: helpers assume compatible tiling, but layouts are incompatible

2. Manual Copy Attempts

Attempt 1: Get thread coordinates with cute.coord(tStS0) → AttributeError: module 'cutlass.cute' has no attribute 'coord'

Attempt 2: Simple test pattern with dynamic indexing → DSLRuntimeError: object cannot be interpreted as an integer

• CuTeDSL JIT requires compile-time constants or vectorized loops
• Dynamic SMEM offset computation not trivial

Key Unknowns Blocking Progress

1. Coordinate Systems: How to get thread's logical position in QK C-fragment partition?
2. Mapping Formula: Given QK C-fragment coordinate (i,j), what is corresponding PV A-operand coordinate?
3. Offset Computation: How to compute SMEM offset from logical coordinate given a layout?
4. Manual Pattern: What's the CuTeDSL pattern for manual layout conversion?

Hypotheses

Hypothesis A: Layouts represent same 128×128 data

• QK C-fragment: tiles for QK MMA (C-fragment partitioning)
• PV A-operand: tiles for PV MMA (A-operand partitioning)
• Need coordinate transformation: QK tile → PV tile mapping

Hypothesis B: tStS0 is whole 128×128 matrix view

• Not thread's slice but entire matrix in TMEM
• Each thread accesses portion via its thread coordinates
• Need thread's position within QK C-fragment partition

Hypothesis C: Manual transpose required

• Implement 128×128 transpose between two tiling patterns
• Could brute-force with nested loops over 128×128
• But thread owns only subset (partitioned by QK MMA)

Questions for CUTLASS LLM

1. Thread Coordinates: How to get thread's logical coordinates within a layout partition?

   • cute.coord() doesn't exist
   • Need: thread position in QK C-fragment partition
   • Possibly: cute.arch.thread_idx(), cute.arch.lane_idx(), or from MMA?

2. Layout Mapping: Given source layout L1 and destination layout L2 for same logical space (128×128 matrix), how to compute coordinate transform?

   • L1: ((128,128),1,1):((65536,1),0,0) (QK C-fragment)
   • L2: ((128,16),1,(4,2),1):((64,1),0,(16,8192),0) (PV A-operand)
   • Need mapping: L1 coordinate → L2 coordinate

3. Offset API: Is there cute.offset(layout, coord)? How to compute physical offset from logical coordinate given a layout?

4. Manual Pattern: Example of manual copy between incompatible layouts in CuTeDSL?

   • Source: tensor with layout L1
   • Destination: tensor with layout L2
   • Manual element-wise copy with address computation

5. MMA Partitioning: How does QK MMA partition the 128×128 C-fragment among threads?

   • Each thread owns subset of elements
   • Need to know which elements thread owns

Found Clue in vendored Code

In moe_torch_scaled_grouped_mm.py:

bidx, bidy, bidz = cute.arch.block_idx()
mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)

• Uses block index to compute MMA tile coordinate
• tiled_mma.thr_id.shape gives thread partition shape
• But we need thread coordinate within warp/CTA, not block

Current Code State

In fmha.py, SMEM-P section is:

else:
    # SMEM-P: Manual addressing (helpers are a trap)
    print(f"[SMEM-P MANUAL] Starting manual P write to SMEM")
    print(f"[SMEM-P MANUAL] tStS0 layout: {tStS0.layout}")
    print(f"[SMEM-P MANUAL] sP layout: {sP.layout}")
    
    # TODO: Implement proper mapping
    # For now, zero sP (wrong but compiles)
    for j in cutlass.range(cute.size(sP), vectorize=True):
        sP[j] = BFloat16(0.0)

Test Status

• test_fmha_v3_stage_d1.py fails with cosine NaN for hd=64, -0.047 for hd=256, 0.011 for hd=512
• Zeroed sP causes PV MMA to read zeros → garbage output
• TMEM-P path works for hd=64 but fails for hd>64 due to TMEM layout mismatch

Progress Update (2026-05-23 19:35 UTC)

CUTLASS LLM responded! Got complete solution:

Key Solution:

1. Use make_identity_tensor(tStS0.shape) for coordinate tensor
2. Partition coordinate tensor same way as data tensor
3. Mapping formula: QK ((m, n), 0, 0) → PV ((m, n % 16), 0, ((n // 16) % 4, n // 64), 0)
4. Use tensor indexing sP[dst_coord] = value, not manual offsets

Implementation Progress:

• ✅ Coordinate mapping function implemented and works (qk_to_pv_coord)
• ✅ Tensor indexing with coordinate works (sP[test_coord] = value)
• ❌ Need to implement full 128-value mapping per thread
• ❌ Need to get QK coordinates for each of thread's 128 P values

Implementation Status (2026-05-23 19:55 UTC)

✅ Implemented full SMEM-P with coordinate mapping

• Created coordinate tensor cS (already existed for row_max)
• Partitioned as tTMEM_LOADcS_frg matching P value fragments
• In softmax loop, for each (k,j):
  • Get QK coordinate (m,n) from tTMEM_LOADcS_frg[k,j]
  • Map to PV SMEM coordinate using formula
  • Write P value (or test pattern) to sP[pv_coord]

❌ Result: Cosine ~0.02 (near zero correlation)

• Kernel compiles and runs
• PV reads SOMETHING from SMEM (output non-zero)
• But mapping appears wrong (random correlation)
• Output scaling huge (280k vs reference 0.2)

Possible Issues:

1. Coordinate mapping formula wrong — PV A-operand layout might differ
2. SMEM swizzle mismatch — tensor indexing might not handle swizzle correctly
3. Thread collisions — Multiple threads writing same SMEM location
4. P value normalization — Unnormalized P values cause scaling issues

Debug Attempts:

1. Test pattern (k+j)*0.01 → cosine 0.02
2. Linear index m*128+n → cosine 0.006 (huge output as expected)
3. Both show mapping is bijective but wrong locations

Next Actions:

1. Verify coordinate mapping by computing SMEM offset manually
2. Check if PV expects transposed P matrix
3. Examine PV MMA tiler and SMEM layout generation
4. Consider alternative: fix TMEM layout generation instead

Progress Update (2026-05-23 21:30 UTC)

TMEM-P path now works at hd=64 (cos 0.973). The root cause of NaN/zeros was a missing TMEM column offset on tOrP0 — PV MMA was reading from column 0 (where S is) instead of column 32 (where P is stored by softmax warps). Fixed with const_expr conditional.

SMEM-P remains unsolved. The make_tiled_copy_C approach gives rank mismatch. Manual coordinate mapping compiles but produces near-zero cosine (wrong addresses). The CUTLASS reference FMHA uses TMEM-P exclusively (12-warp layout with more TMEM budget). For our 6-warp layout, SMEM-P is needed for hd>64.

Current D1 status:

• hd=64 (TMEM-P): cos 0.973 ✅
• hd=256 (SMEM-P stub): FAIL (zeros)
• hd=512 (SMEM-P stub): FAIL (zeros)

Workaround for hd>64: The D5b milestone (Python SWA+sink merge) works at hd=64. SMEM-P for hd>64 is a production optimization, not a correctness blocker. The full DSV4 pipeline (CSA + HCA + SWA) can be tested at hd=64 with TMEM-P.

Key discoveries:

1. make_tiled_copy_C(store_atom, qk_mma) creates a copy that partitions threads by QK C-fragment layout, but the source and destination have incompatible ranks (4 vs 3). This is a fundamental layout incompatibility.
2. Manual coordinate mapping (qk_to_pv_coord) compiles and runs but produces wrong results. The mapping formula may be incorrect, or SMEM swizzle may interfere with tensor indexing.
3. The CUTLASS reference FMHA (12-warp) avoids SMEM-P entirely by using TMEM-P with more warps (more TMEM budget). A 12-warp layout would solve the SMEM-P problem architecturally.

Recommended next steps for SMEM-P:

1. Try 12-warp layout (like CUTLASS reference) to avoid SMEM-P entirely
2. OR: Use make_tiled_copy with pv_mma (not qk_mma) for the copy, since PV MMA knows the SMEM layout
3. OR: Implement a two-stage copy: QK C-fragment → intermediate buffer → PV A-operand SMEM