[Feature] Add GVA support for Lightning#85
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces Grouped Value Attention (GVA) support to the lightning attention SM100 operators, allowing the number of value heads (HV) to be greater than the query/key heads (H) as long as it is divisible by H. The changes span the core operator implementation, the benchmark suite, and the test suite. The review feedback identifies critical safety improvements for the CUDA kernels, specifically recommending that input tensors (Q, K, V, initial_state, and state_pool) be made contiguous to prevent silent memory corruption, and that the state_pool size be validated to avoid out-of-bounds memory accesses during variable-length sequence processing.
|
@sjmshsh the conflicts need to be resolved, PTAL |
# Conflicts: # tests/test_la_decode.py
done |
|
solved @icavan |
#81
python tests/test_lightning_attn.py --test all
python tests/test_lightning_attn.py --test ref
python tests/test_lightning_attn.py --test fla
python tests/test_lightning_attn.py --test h0ht
python tests/test_lightning_attn.py --test varlen
python -m pytest tests/test_la_decode.py tests/test_la_decode_pool.py -v -s
python -m pytest tests/test_la_decode.py tests/test_la_decode_pool.py -v -s -k "gva or prefill_decode_e2e"
@icavan @KevinZeng08