Support variable-length sequences for mamba block with position indices

[ptxu78]

Enable the mamba block to support variable-length sequence inputs using positional encoding. Passing Positional Indices results in negligible performance loss for the mamba block. For common variable-length sequence distributions, performance can be improved by 4-6x.

• Modification:
  Refer to the link [Feature] Support variable-length sequences for mamba block #244 to replace cumulative sequences with a positional encoding matrix. The position encoding is more suitable for parallel acceleration and is more commonly found in the outputs of various dataloaders.

For example, a packaged sequence of length 16 consists of four independent sub-sequences with lengths of 3, 5, 6, 2. Then corresponding:
Cumulative sequence: [0, 3, 8, 14, 16]
Position encoding: [0, 1, 2, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1]

In the Mamba module, there are two steps that are not sequence-wise: conv1d and selective_scan. Sub-sequences within the same sequence can affect each other, and we have modified causal-conv1d and selective_scan. These two CUDA operators are implemented using position encoding to eliminate the mutual influence between sub-sequences.

• How to use:

1. Setup two cuda_operators:

git clone --branch feat/pack_with_position_indices https://github.com/ptxu78/pack_mamba.git
python setup .../pack_mamba/setup_onlyCUDA.py install
git clone --branch feat/pack_with_position_indices https://github.com/ptxu78/causal-conv1d-pack.git
python setup .../causal-conv1d-pack/setup_onlyCUDA.py install

2. Pack the variable-length sequences and input them, along with the corresponding position encoding.

[ScottHoang]

This PR attempts to resolve the issue derived from training with packed data?

[ptxu78]

This PR attempts to resolve the issue derived from training with packed data?

Yes, this PR allow Mamba to handle packed data more effectively: it significantly increases throughput while ensuring the mathematical equivalence of the training results.