Integrate Flex Decoding

[BoyuanFeng]

This PR integrates flex decoding with gpt-fast.

End-to-end performance gain of Llama2-7b

Device: H100
Unit: tokens/sec

Length	spda	Flex Decoding	Speedup
1024	143.57	142.59	0.99x
2048	138.67	140.13	1.01x
4096	128.85	135.31	1.05x
8192	111.59	125.38	1.12x
16384	89	109.8	1.23x
32768	64.11	87.67	1.37x

command:

export MODEL_REPO=meta-llama/Llama-2-7b-chat-hf
python generate.py --compile --max_new_tokens 16384 --checkpoint_path checkpoints/$MODEL_REPO/model.pth --prompt "Hello, my name is"

Please also set ModelArgs.block_size = 65536 to repeat the result.

We expect to see larger speedup on longer context length.

[Chillee]

get_mask_mod shouldn't take an integer - it should take a mask_mod. We also don't need to set at as an argument, just set it as an attribute within the module.

[Chillee]

Specifically, you should be able to take any existing mask_mod and wrap it to make it automatically support an offset.

[Chillee]

Mostly looks good (modulo some other nits), although we'll probably want to wait on landing this until the other PRs on pytorch core are landed at least.

[merrymercy]

This looks interesting. I would like to share some numbers we got with torch.compile + flashinfer in sglang. It can serve as some good baselines. To run the 32k one, you need to edit the config.json with "max_position_embeddings": 65536,.

# length = 1024
# Decode.  median latency: 0.00596 s, median throughput:    167.67 token/s
python3 -m sglang.bench_latency --model meta-llama/Llama-2-7b-chat-hf --batch-size 1 --input 1024 --output 8 --enable-torch-compile

# length = 32768
# Decode.  median latency: 0.01136 s, median throughput:     88.01 token/s
python3 -m sglang.bench_latency --model meta-llama/Llama-2-7b-chat-hf --batch-size 1 --input 32768 --output 8 --enable-torch-compile

You can find more numbers at sgl-project/sglang#1008

[Chillee]

@merrymercy We run on nerfed H100s internally at Meta with only 2.4 TB/s of bandwidth, so these numbers aren't 1:1 comparable.

But it's a good comparison :)

[BoyuanFeng]

offline discussed that BlockMask getitem sets mask_mod as None and the user needs to specify the correct mask_mod. In GPT-Fast, we rely on model.get_mask_mod to do so.

[drisspg]

try doing
create_block_mask_compile = torch.compile(create_block_mask)

as a global