[Feature] Refactor the RPA code that’s limited by SMEM size so it can support longer contexts #480
Description
Checklist
- 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
- 2. Please use English, otherwise it will be closed.
Motivation
Because of SMEM size limits, our support for long contexts is currently constrained. We need to refactor the related parameter-passing logic in RPA to support longer contexts.
Related resources
No response