gpt_bigcode: added FusedSDPA kernel

[mgonchar]

Added support of following options to gpt_bigcode (starcoderbase) model use_flash_attention,
flash_attention_recompute,
flash_attention_fast_softmax,
flash_attention_causal_mask

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[mgonchar]

Original implementation

python run_generation.py \
    --model_name_or_path bigcode/starcoderbase-3b \
    --use_hpu_graphs \
    --use_kv_cache \
    --batch_size 42 \
    --max_new_tokens 1024 \
    --do_sample\
    --bf16

Stats

---------------------------------------------------------------------------------------------------------------
Throughput (including tokenization) = 3970.2326582228557 tokens/second
Number of HPU graphs                = 15
Memory allocated                    = 84.78 GB
Max memory allocated                = 94.62 GB
Total memory available              = 94.62 GB
Graph compilation duration          = 64.8335093039932 seconds
---------------------------------------------------------------------------------------------------------------

FusedSDPA

python run_generation.py \
    --model_name_or_path bigcode/starcoderbase-3b \
    --use_hpu_graphs \
    --use_kv_cache \
    --batch_size 42 \
    --max_new_tokens 1024 \
    --do_sample\
    --bf16 \
    --use_flash_attention

Stats

--------------------------------------------------------------------------------------------------------------
Throughput (including tokenization) = 4175.830644970393 tokens/second
Number of HPU graphs                = 15
Memory allocated                    = 13.41 GB
Max memory allocated                = 94.45 GB
Total memory available              = 94.62 GB
Graph compilation duration          = 75.58716204599841 seconds
--------------------------------------------------------------------------------------------------------------

[vidyasiv]

@mgonchar , thank you for your PR.

• Please run make style prior to future submissions, it takes care of code formatting fixes
• Please verify if bigcode/starcoder which also uses this model file runs alright for you. I am seeing an error RuntimeError: Graph compile failed. synStatus=synStatus 26 [Generic failure]. for python run_generation.py --model_name_or_path bigcode/starcoder --batch_size 1 --use_hpu_graphs --use_kv_cache --max_new_tokens 100 --bf16 --use_flash_attention. I suspect this has to do with "starcoderbase" check and it works if substituted with "starcoder" but we need a better match because we want to avoid starcoder2 perhaps?
• Testing: Please add a test for starcoder, starcoderbase with flash attention options in tests/test_text_generation_example.py

I will do a another pass after we resolve these

[mgonchar]

@vidyasiv I've updated this PR, based on your feedback. Please have a look

[vidyasiv]

Minor typo.I was able to run tests and so far LGTM

[vidyasiv]

Please re-run make style

[mgonchar]

@vidyasiv done

[vidyasiv]

@regisss , please take a look

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[mgonchar]

PR rebased. I've rechecked rebased code - no regressions found. @regisss please have a look

[regisss]

LGTM!