Skip to content

[AMD] Fix AITER unified decode page table for page size one - #40410

Open
vorapolsiloai wants to merge 1 commit into
sgl-project:mainfrom
vorapolsiloai:fix/aiter-page1-cuda-graph-table
Open

vorapolsiloai wants to merge 1 commit into
sgl-project:mainfrom
vorapolsiloai:fix/aiter-page1-cuda-graph-table

Conversation

@vorapolsiloai

@vorapolsiloai vorapolsiloai commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

With AITER unified attention (SGLANG_USE_AITER_UNIFIED_ATTN=1), page size 1 (the default), and CUDA graphs enabled, decode reads a page table that was never filled in, so the output is wrong.

Before each CUDA graph replay, AiterAttnBackend._apply_cuda_graph_metadata() copies the request-to-token indices into the persistent cuda_graph_page_table. It did this for the sliding-window KV pool and for page_size > 1. It skipped page_size == 1, so the replay read the table's initial or stale contents. Eager decode (init_forward_metadata) builds this table for every page size, so only graph replay is affected. Query and KV dtypes do not matter.

I found this while validating #39998. The bug is in the existing decode metadata path, so the fix is a separate PR.

Modifications

  • In the CUDA graph decode metadata path, always copy the page indices into cuda_graph_page_table. Convert them with _transform_table_1_to_real only when page_size > 1. This matches the eager path.
  • Add test_cuda_graph_decode_copies_page_size_one_table. It calls _apply_cuda_graph_metadata() with page_size=1 and checks that the graph page table holds the req_to_token rows.

Accuracy Tests

Speed Tests and Profiling

This is a correctness fix. It adds one device-to-device copy per graph replay, the same copy that page sizes above 1 already do. It is not meant to change performance.

Checklist

  • Format code with pre-commit.
  • Add unit tests.
  • Provide accuracy results.

CI States

Latest PR Test (Base): ❌ Run #37836097801
Latest PR Test (Extra): ❌ Run #37836097439
Latest PR Test (AMD ROCm 10): ❌ Run #37836097712

Unified decode reuses a capture-stable page table during CUDA graph replay, but the page-size-one path never populated it. Copy the request-to-token indices into that table just as the paged paths do.

- Populate generic unified-decode metadata for page size one
- Add a regression that fails with the zero-initialized graph table

This branch has not been deployed

No deployments
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