mark_hardware_transpose_candidates breaks elements-per-thread invariant
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
After running `mark_elements_per_thread` on attention with the following config:
```
attention, hyperparams, _ = get_vanilla_attention_kernel(
AttentionShape(
num_query_heads=8,
num_kv_heads=2,
query_seq_len=256,
head_size_kv=64,
head_size=64,
kv_seq_len=256,
),
(MMAType.F32_16x16x16_F16, MMAType.F32_16x16x16_F16),
False,
)
```
reads from shared memory have `elements_per_thread = 4` and `index={B: $WG2*BLOCK_B : 1 : 1, N: $WG1*BLOCK_N + 4*(Mod($T0, 4)) : 1 : 1, K2: $ARGK2*BLOCK_K2 + floor((Mod($T0, 64))/4) : 8 : 1})`, the latter implying `elements_per_thread = 8`. Note that prior to that pass, the index uses `size=4` that _is_ consistent. It is unclear if downstream passes in the pipeline take one or the other value as source-of-truth.
This yet again highlights the blatant lack of pass-level testing and verifiers and over-reliance on e2e testing that happen to work in narrow cases.
Contributor guide
Assessment
This issue has not been assessed yet.