Motivation

Currently causal_conv1d_fn triton kernel is used for extend & draft_extend. From the profile, it is inefficient - many CPU<>GPU copy sync and its grid method is very slow (takes 139us to calculate the triton grid) while the real triton kernel launch itself only takes 17us.

Modifications

Take this request as example:

• seq_lens = [1, 9, 2, 17, 4]
• BLOCK_M=8
• num_of_blockm_s_per_seq_len = torch.tensor([1, 2, 1, 3, 1]) -> total_blockm_s = 1+2+1+3+1=8
• m_list = [0, 1, 1, 2, 3, 3, 3, 4]: len(m_list) = 8 = total_blockm_s
• offset_list = [0, 0, 1, 0, 0, 1, 2, 0]: len(offset_list) = 8 = total_blockm_s

Before

• Each program_id(0) in _causal_conv1d_fwd_kernel processes one BLOCK_M out of total 8 BLOCK_M
• Each BLOCK_M maps to one block of tokens. program_id(0) = i maps to blockm[i], maps to (m_list[i], offset_list[i]) chunk of tokens
• For seq_lens = [1, 9, 2, 17, 4] & BLOCK_M=8, we have the following programs (aligns with m_list = [0, 1, 1, 2, 3, 3, 3, 4] and offset_list = [0, 0, 1, 0, 0, 1, 2, 0]):
  -- program_id(0) = 0
  -- program_id(0) = 1
  -- program_id(0) = 2
  -- program_id(0) = 3
  -- program_id(0) = 4
  -- program_id(0) = 5
  -- program_id(0) = 6
  -- program_id(0) = 7

After: same processing logic as in triton_backend.py's extend_attention() method

• Each program_id(0) in _causal_conv1d_fwd_kernel processes one seq_len out of total 5 seq_lens
• Each program_id(1) in _causal_conv1d_fwd_kernel processes one block of tokens belonging to that seq_len
• For seq_lens = [1, 9, 2, 17, 4] & BLOCK_M=8, we have the following programs (aligns with num_of_blockm_s_per_seq_len = torch.tensor([1, 2, 1, 3, 1])):
  -- program_id(0) = 0, program_id(1) = 0
  -- program_id(0) = 1, program_id(1) = 0
  -- program_id(0) = 1, program_id(1) = 1
  -- program_id(0) = 2, program_id(1) = 0
  -- program_id(0) = 3, program_id(1) = 0
  -- program_id(0) = 3, program_id(1) = 1
  -- program_id(0) = 3, program_id(1) = 2
  -- program_id(0) = 4, program_id(1) = 0

By doing so, we don't have to calculate m_list & offset_list which requires cpu<> gpu copy and expensive tensor/numpy operations. Instead, we can just leverage forward_batch.extend_seq_len_cpu as it is in the triton kernel.

In addition to this main change, there are some misc changes done by this PR such as removing unnecessary code/tensor and improve the if condition in triton kernel

Benchmarking and Profiling

causal_conv1d_fn runtime

Before: p50=313us; avg=656us

After: p50=95us; avg=108us - 3X faster

GSM8k Accuracy

GSM8k Output token/s

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.