Lightning-attn supports pretransposed states

[icavan]

[higgsboson1710]

"Hi @icavan! I'm a CS student interested in Machine Learning and I'd love to work on this as my first contribution to cuLA. Since there isn't a description yet, could you point me toward the files that need to be updated to support pretransposed states? Also, could you please assign this to me?"

[icavan]

"Hi @icavan! I'm a CS student interested in Machine Learning and I'd love to work on this as my first contribution to cuLA. Since there isn't a description yet, could you point me toward the files that need to be updated to support pretransposed states? Also, could you please assign this to me?"

@higgsboson1710 you could refer to #33, which is a KDA pretranspose implementation.

[icavan]

"Hi @icavan! I'm a CS student interested in Machine Learning and I'd love to work on this as my first contribution to cuLA. Since there isn't a description yet, could you point me toward the files that need to be updated to support pretransposed states? Also, could you please assign this to me?"

@higgsboson1710 Below lists some suggestions from AI. I believe it is good enough to start with.

Background

Lightning Prefill currently reads/writes state in BHKV [B, H, D_k, D_v] layout, while LA Decode expects BHVK [B, H, D_v, D_k]. The goal is to align prefill state I/O to BHVK, eliminating the transpose between prefill and decode.

Task 1: Modify fstate_layout stride definition

• cula/ops/lightning_attn.py:413-416
• Change stride from (1, D, ...) (K-inner) to (D, 1, ...) (V-inner)

Task 2: Modify h0 load path (initial state input)

• cula/ops/lightning_attn.py:1705-1707
• Thread local_tidx currently reads row local_tidx; change to read column local_tidx
• Evaluate whether an SMEM transpose buffer is needed for coalesced access

Task 3: Modify ht store path (final state output)

• cula/ops/lightning_attn.py:1880-1883
• Symmetric change: write column instead of row
• The varlen inplace update path writes back to the initial_state pool via the same code — must be updated as well

Task 4: Update Python API state allocation

• lightning_attn_fwd (cula/ops/lightning_attn.py:2919-2920) — adjust ht allocation stride convention
• lightning_attn_fwd_varlen (cula/ops/lightning_attn.py:3085) — adjust state_pool allocation
• Document that callers must pass initial_state in BHVK format

Task 5: Evaluate GMEM access performance

• Column-wise per-thread access may cause uncoalesced stores
• Evaluate whether an SMEM transpose buffer is needed (reference: existing StMatrix/LdMatrix transpose in the K weighting path)
• Alternative: zero-copy logical transpose via swapped CuTe strides at the GMEM interface

Task 6: Update tests

• tests/test_lightning_attn.py — create state tensors in BHVK layout
• tests/test_la_decode.py — remove the manual BHKV→BHVK transpose in run_la_decode(), since prefill output will already be BHVK
• Add an end-to-end prefill → decode test verifying state passes directly without transpose