Repository navigation
AITER MLA backend: Avoid CPU sync in _build_decode - #35765
Conversation
This uses the same triton kernel as flashinfer.py to prevent a CPU sync when preparing the page_indices buffer. Signed-off-by: Patrick Schlangen <pschlan@amd.com>
There was a problem hiding this comment.
Code Review
This pull request successfully resolves a CPU synchronization issue in the ROCm AITer MLA backend by replacing a data-dependent indexing operation with a custom Triton kernel. The approach is sound and aligns with the PR's objective. I've identified a potential critical issue in the new Triton kernel related to memory access patterns, which could lead to incorrect behavior if the block_table_tensor is not contiguous on its last dimension. My suggestion makes the kernel more robust by explicitly handling tensor strides.
| @triton.jit | ||
| def _copy_page_indices_kernel( | ||
| page_indices, | ||
| block_table, | ||
| block_table_stride, | ||
| cu_num_blocks, | ||
| BLOCK_SIZE: tl.constexpr, | ||
| ): | ||
| """Copy block table rows into a flat page_indices buffer using indptr. | ||
| Avoids blocking boolean mask indexing (tensor[mask]) which has | ||
| data-dependent output size and forces sync. | ||
| This is the same kernel as introduced in backends/flashinfer.py. | ||
| """ | ||
| req_idx = tl.program_id(0) | ||
| row_ptr = block_table + req_idx * block_table_stride | ||
| start_idx = tl.load(cu_num_blocks + req_idx) | ||
| end_idx = tl.load(cu_num_blocks + req_idx + 1) | ||
| num_blocks = end_idx - start_idx | ||
|
|
||
| offset = tl.arange(0, BLOCK_SIZE) | ||
| for i in tl.range(0, num_blocks, BLOCK_SIZE): | ||
| block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks) | ||
| tl.store( | ||
| page_indices + start_idx + i + offset, | ||
| block_ids, | ||
| mask=i + offset < num_blocks, | ||
| ) |
There was a problem hiding this comment.
The new Triton kernel _copy_page_indices_kernel assumes that the block_table tensor is contiguous along its last dimension (i.e., stride(1) == 1) because it uses simple pointer arithmetic (row_ptr + i + offset) to access elements within a row. If block_table_tensor is ever a non-contiguous view (e.g., from transpose or advanced indexing), this kernel will read incorrect data, leading to silent correctness issues or crashes.
To make the kernel more robust, it's better to explicitly use the stride of the second dimension for memory access. This ensures correctness regardless of the tensor's memory layout.
@triton.jit
def _copy_page_indices_kernel(
page_indices,
block_table,
block_table_stride_0,
block_table_stride_1,
cu_num_blocks,
BLOCK_SIZE: tl.constexpr,
):
"""Copy block table rows into a flat page_indices buffer using indptr.
Avoids blocking boolean mask indexing (tensor[mask]) which has
data-dependent output size and forces sync.
This is the same kernel as introduced in backends/flashinfer.py.
"""
req_idx = tl.program_id(0)
row_ptr = block_table + req_idx * block_table_stride_0
start_idx = tl.load(cu_num_blocks + req_idx)
end_idx = tl.load(cu_num_blocks + req_idx + 1)
num_blocks = end_idx - start_idx
offset = tl.arange(0, BLOCK_SIZE)
for i in tl.range(0, num_blocks, BLOCK_SIZE):
current_offsets = i + offset
mask = current_offsets < num_blocks
block_ids = tl.load(row_ptr + current_offsets * block_table_stride_1, mask=mask)
tl.store(
page_indices + start_idx + current_offsets,
block_ids,
mask=mask,
)There was a problem hiding this comment.
This doesn't seem to be a concern in flashinfer (where this kernel is invented), but will try to check the codebase if this can become an issue
| _copy_page_indices_kernel[(num_reqs,)]( | ||
| self.paged_kv_indices, | ||
| block_table_tensor, | ||
| block_table_tensor.stride(0), | ||
| paged_kv_indptr, | ||
| BLOCK_SIZE=1024, | ||
| ) |
There was a problem hiding this comment.
To support the suggested change in the _copy_page_indices_kernel, you'll need to pass block_table_tensor.stride(1) to the kernel as well.
| _copy_page_indices_kernel[(num_reqs,)]( | |
| self.paged_kv_indices, | |
| block_table_tensor, | |
| block_table_tensor.stride(0), | |
| paged_kv_indptr, | |
| BLOCK_SIZE=1024, | |
| ) | |
| _copy_page_indices_kernel[(num_reqs,)]( | |
| self.paged_kv_indices, | |
| block_table_tensor, | |
| block_table_tensor.stride(0), | |
| block_table_tensor.stride(1), | |
| paged_kv_indptr, | |
| BLOCK_SIZE=1024, | |
| ) |
|
@pschlan-amd I think it looks good. Before approving can you share the lm_eval score of DeepSeek-R1 model after this changes? and also share end-to-end benchmark metrics to show the perf gain? |
Thanks for checking! I ran lm_eval / gsm8k against DeepSeek-R1-0528, TP=8 on a test system with 8x MI355X. Pre change Post change Latency benchmark results (vllm bench latency) with the same model: Pre change Post change |
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
This uses the same triton kernel as flashinfer.py to prevent a CPU sync when preparing the page_indices buffer.
NOTE: This is my first time working on this type of code, and I'd appreciate thorough reviews.
Purpose
Fix a long sync when running _build_decode. This was found during profiling.
Before the change:
After the change:
(not up to scale)
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.