linkedin / linkedin/Liger-Kernel
Feature Request: Add THD / Packed-Sequence Training Support
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
### 🚀 The feature, motivation and pitch
## Motivation
Some training frameworks represent variable-length batches in a packed THD layout, where activations are shaped like `(total_tokens, num_heads, head_dim)`, or more generally `(total_tokens, hidden_dim)`, instead of the standard padded batch layout `(batch_size, seq_len, ...)`.
This layout is useful for sequence packing and variable-length training because it avoids padding overhead.
Supporting this layout may require passing sequence metadata such as `cu_seqlens` or `max_seqlen` to kernels that need per-sequence position information.
## Request
Add support for THD / packed-sequence training layouts across Liger kernels.
The long-term goal is to support inputs such as `hidden_states: (total_tokens, hidden_dim)` and attention tensors such as `(total_tokens, num_heads, head_dim)`, with metadata such as `cu_seqlens` and `max_seqlen` where needed.
### Alternatives
_No response_
### Additional context
_No response_
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
The issue names no implementation files, entry points, or tests. Start by surveying the existing Liger kernels and their shape-related tests, then determine which kernels need THD inputs and sequence metadata such as cu_seqlens and max_seqlen. Done means the supported kernels accept packed layouts without padding assumptions and have coverage for variable-length training.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 28/100