Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Kernel][Model] Varlen prefill + Prefill chunking support for mamba kernels and Jamba model #8533

Merged
merged 53 commits into from
Sep 29, 2024
Merged

[Kernel][Model] Varlen prefill + Prefill chunking support for mamba kernels and Jamba model #8533

merged 53 commits into from
Sep 29, 2024

Commits on Sep 15, 2024

  1. reduce x size from float2 to float

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    b8e4cc6 View commit details
    Browse the repository at this point in the history

  2. Support single chunk as input

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    8991183 View commit details
    Browse the repository at this point in the history

  3. working with grid dimensions x = batch size, y = max seqlen, same amount 

    launches, lower memory footprint
    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    ea0089f View commit details
    Browse the repository at this point in the history

  4. final and initial states suppport

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    cf60b69 View commit details
    Browse the repository at this point in the history

  5. working with cache indices

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    3998748 View commit details
    Browse the repository at this point in the history

  6. WIP - add varlen to ssm

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    4038d84 View commit details
    Browse the repository at this point in the history

  7. Working version with splits , TBD clean up

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    949718f View commit details
    Browse the repository at this point in the history

  8. also working with init prefill

    @mzusman
    mzusman committed Sep 15, 2024
    Configuration menu
    Copy the full SHA
    3e90085 View commit details
    Browse the repository at this point in the history

Commits on Sep 16, 2024

  1. Remove last channel kernel

    @mzusman
    mzusman committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    b01f705 View commit details
    Browse the repository at this point in the history

  2. Clean up causal_conv1d

    @mzusman
    mzusman committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    b9144f2 View commit details
    Browse the repository at this point in the history

  3. Clean up selective_scan kernels and torch bindings

    @mzusman
    mzusman committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    d2b97fa View commit details
    Browse the repository at this point in the history

  4. fix tests

    @mzusman
    mzusman committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    869aaf1 View commit details
    Browse the repository at this point in the history

  5. Fix wrappers

    @mzusman
    mzusman committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    0addc82 View commit details
    Browse the repository at this point in the history

Commits on Sep 17, 2024

  1. take off requirement for stride -1 == 1

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    c4fe338 View commit details
    Browse the repository at this point in the history

  2. Update causal_conv1d_update to use the new kernel

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    d6fe5fd View commit details
    Browse the repository at this point in the history

  3. ssm state to be able to use different dtypes (itype)

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    f2411b3 View commit details
    Browse the repository at this point in the history

  4. more causal_conv1d_update fixes

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    216462f View commit details
    Browse the repository at this point in the history

  5. Add prefill chunking to jamba modeling file

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    3c6ec5c View commit details
    Browse the repository at this point in the history

  6. Fix formating

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    4753054 View commit details
    Browse the repository at this point in the history

  7. remove print

    @mzusman
    mzusman committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    1581443 View commit details
    Browse the repository at this point in the history

Commits on Sep 22, 2024

  1. remove cruft and add comments

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    1e08a4e View commit details
    Browse the repository at this point in the history

  2. Add guards

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    674e9f9 View commit details
    Browse the repository at this point in the history

  3. Renaming and fix bug for short sequences

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    1b9b3ba View commit details
    Browse the repository at this point in the history

  4. renaming and add test for random cache indices and random has initial… 

    … states
    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    8e4f92d View commit details
    Browse the repository at this point in the history

  5. Add comments

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    3a8632d View commit details
    Browse the repository at this point in the history

  6. Add comment on jamba tests

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    6a9acb7 View commit details
    Browse the repository at this point in the history

  7. has initial state as bool and add comments to jamba

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    9408615 View commit details
    Browse the repository at this point in the history

  8. Merge remote-tracking branch 'github/main' into varlen_mamba_causal_c… 

    …onv1d_upstream
    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    254351e View commit details
    Browse the repository at this point in the history

  9. Format

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    d3d4e0f View commit details
    Browse the repository at this point in the history

  10. Some alignments with the changed from upstream

    @mzusman
    mzusman committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    801cd7a View commit details
    Browse the repository at this point in the history

Commits on Sep 25, 2024

  1. Remove cruft

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    ddf1d5c View commit details
    Browse the repository at this point in the history

  2. Fix prefill chunking test

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    8278263 View commit details
    Browse the repository at this point in the history

  3. Remove unused returns

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    fbd1756 View commit details
    Browse the repository at this point in the history

  4. Use varlen in Jamba

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    2153e03 View commit details
    Browse the repository at this point in the history

  5. Use decode kernels

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    64c2f4b View commit details
    Browse the repository at this point in the history

  6. Remove comment

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    b4515f7 View commit details
    Browse the repository at this point in the history

  7. Merge remote-tracking branch 'github/main' into varlen_mamba_causal_c… 

    …onv1d_upstream
    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    693e818 View commit details
    Browse the repository at this point in the history

  8. Fix opcheck for mamba ssm and causal conv1d

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    82b3a2a View commit details
    Browse the repository at this point in the history

  9. Put back the figure

    @mzusman
    mzusman committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    fee189d View commit details
    Browse the repository at this point in the history

Commits on Sep 26, 2024

  1. WIP - fix opcheck tests

    @mzusman
    mzusman committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    d74ee9c View commit details
    Browse the repository at this point in the history

  2. Formating and fix opcheck tests

    @mzusman
    mzusman committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    d4ddb12 View commit details
    Browse the repository at this point in the history

  3. Add final state out the selective_scan_ref could fix tests fail

    @mzusman
    mzusman committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    5528cfa View commit details
    Browse the repository at this point in the history

Commits on Sep 27, 2024

  1. Fix test failures

    @mzusman
    mzusman committed Sep 27, 2024
    Configuration menu
    Copy the full SHA
    9fb3353 View commit details
    Browse the repository at this point in the history

  2. Merge remote-tracking branch 'github/main' into varlen_mamba_causal_c… 

    …onv1d_upstream
    @mzusman
    mzusman committed Sep 27, 2024
    Configuration menu
    Copy the full SHA
    f9bc2fa View commit details
    Browse the repository at this point in the history

  3. format

    @mzusman
    mzusman committed Sep 27, 2024
    Configuration menu
    Copy the full SHA
    9c6d140 View commit details
    Browse the repository at this point in the history

Commits on Sep 29, 2024

  1. renaming and sort out the set_params functions

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    1362e9b View commit details
    Browse the repository at this point in the history

  2. add comment on final state assigment

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    b8580f5 View commit details
    Browse the repository at this point in the history

  3. renaming

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    209a6a9 View commit details
    Browse the repository at this point in the history

  4. renaming

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    c670e6d View commit details
    Browse the repository at this point in the history

  5. Renaming

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    30e7239 View commit details
    Browse the repository at this point in the history

  6. Jamba adaptations

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    9e7ecf2 View commit details
    Browse the repository at this point in the history

  7. Fix jamba calls

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    893fdf9 View commit details
    Browse the repository at this point in the history

  8. Formating and renaming

    @mzusman
    mzusman committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    e1c018b View commit details
    Browse the repository at this point in the history