Skip to content

Reduce shared memory usage in decode scheduling metadata kernel - #12

Open
FrankMinions wants to merge 1 commit into
sgl-project:sglfrom
FrankMinions:opt-smem-alloc
Open

Reduce shared memory usage in decode scheduling metadata kernel#12
FrankMinions wants to merge 1 commit into
sgl-project:sglfrom
FrankMinions:opt-smem-alloc

Conversation

@FrankMinions

Copy link
Copy Markdown

Summary

  • Reduce dynamic shared memory in get_mla_metadata_kernel from 5*b+1 to 3*b+1 by dropping redundant first_block_idx / last_block_idx buffers.
  • Since first_token_idx is always 0, first_block_idx is always 0 and last_block_idx + 1 equals num_blocks, so block range bounds can be derived from num_blocks and the local split offset (now_block) instead of being cached in shared memory.
  • Update begin_block_idx / end_block_idx / is_last_req_splitted computation accordingly; scheduling semantics are unchanged.

Motivation

The decode scheduling metadata kernel previously reserved shared memory for five int arrays of size roughly proportional to batch size. Two of those arrays were redundant given the current dense/sparse indexing assumptions (first_token_idx == 0). Removing them lowers shared-memory pressure for large batches and simplifies the kernel. Especially on H200, a large chunk-prefill-size can cause a bottleneck in shared-memory allocation, leading to a fallback to get_mla_metadata_kernel_low_smem.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant