Skip to content

Conversation

@Jeff-Huang
Copy link
Contributor

  • Add multi-dimensional page index support (YsGatherDims) in tile_scatter_gather
  • Add is_gather_dim() and get_gather_index() for multi-dim page lookup
  • Override MakeVDramTileDistribution() for VECTORIZED_LAYOUT to match GEMM's BWarpDstrEncoding (K decomposition: {K2, K0, K1})
  • Add GetGemmKDecomposition() to retrieve kABKLane and kKPerThread
  • Add static_assert for RowMajor VLayout requirement in batch prefill

Proposed changes

Please describe the motivation behind the pull request, whether it enables a new feature or fixes a bug. If there are associated pull requests or issues, please link them to the pull request.

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

  • I have added tests relevant to the introduced functionality, and the unit tests are passing locally
  • I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
  • I have added inline documentation which enables the maintainers with understanding the motivation
  • I have removed the stale documentation which is no longer relevant after this pull request
  • (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
  • I have run clang-format on all changed files
  • Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

- Add multi-dimensional page index support (YsGatherDims) in tile_scatter_gather
- Add is_gather_dim() and get_gather_index() for multi-dim page lookup
- Override MakeVDramTileDistribution() for VECTORIZED_LAYOUT to match
  GEMM's BWarpDstrEncoding (K decomposition: {K2, K0, K1})
- Add GetGemmKDecomposition() to retrieve kABKLane and kKPerThread
- Add static_assert for RowMajor VLayout requirement in batch prefill
@Jeff-Huang Jeff-Huang force-pushed the batch-prefill-improve-vectorized-layout-perf branch from c43b242 to aef6bf2 Compare January 27, 2026 22:40
@Jeff-Huang Jeff-Huang enabled auto-merge (squash) January 28, 2026 07:42
@Jeff-Huang Jeff-Huang disabled auto-merge January 28, 2026 07:42
@poyenc
Copy link
Contributor

poyenc commented Jan 28, 2026

LGTM

@Jeff-Huang Jeff-Huang merged commit e3556fe into develop Jan 28, 2026
23 checks passed
@Jeff-Huang Jeff-Huang deleted the batch-prefill-improve-vectorized-layout-perf branch January 28, 2026 23:18
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.

3 participants