Skip to content

misc: Add runtime validation for plan/run consistency in BatchMLAPagedAttentionWrapper #2246

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bkryu wants to merge 1 commit into from
Closed

misc: Add runtime validation for plan/run consistency in BatchMLAPagedAttentionWrapper #2246

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions flashinfer/mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,6 +346,8 @@ def plan(
self._page_size = page_size
self._sm_scale = sm_scale
self._use_profiler = use_profiler
# Store total query count for validation in run()
self._total_num_queries = int(qo_indptr_host[-1].item())

@yzh119 yzh119 Dec 21, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have some concern if qo_indptr will be a device tensor in the future (because [-1].item() will call a cuda sync in that case). But considering we expect user to pass a host-side tensor at the moment I'm good with this.

@bkryu bkryu Dec 21, 2025 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I initially thought the same, but a few lines above, we have:

        qo_indptr_host = qo_indptr.to("cpu")
        kv_indptr_host = kv_indptr.to("cpu")
        kv_len_arr_host = kv_len_arr.to("cpu")

which means that the plan() function already requires a CUDA sync and is therefore cannot be placed in CUDA graph, so I figured it is okay here


self._plan_info = self._cached_module.plan(
self._float_workspace_buffer,
Expand Down Expand Up @@ -433,6 +435,23 @@ def run(
page_table : Optional[torch.Tensor]
The page table of the paged kv-cache, shape: ``[batch_size, num_pages]``. Required when ``backend`` is ``cutlass``.
"""

if self._backend != "cutlass":

@wenscarl wenscarl Dec 20, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe check for cutlass backend also.

# Validate that run() inputs match the plan() configuration.
# Note: cutlass backend doesn't use plan()
if q_nope.shape[0] != self._total_num_queries:
raise ValueError(
f"q_nope.shape[0] ({q_nope.shape[0]}) does not match "
f"qo_indptr[-1] ({self._total_num_queries}) from plan(). "
f"The total number of query tokens must be consistent between plan() and run()."
)
if q_pe.shape[0] != self._total_num_queries:
raise ValueError(
f"q_pe.shape[0] ({q_pe.shape[0]}) does not match "
f"qo_indptr[-1] ({self._total_num_queries}) from plan(). "
f"The total number of query tokens must be consistent between plan() and run()."
)
Comment on lines +442 to +453

@gemini-code-assist gemini-code-assist Bot Dec 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

These two checks for q_nope and q_pe can be combined to reduce code duplication. A combined check also provides a more comprehensive error message if both tensors have incorrect shapes, which improves the debugging experience.

            if q_nope.shape[0] != self._total_num_queries or q_pe.shape[0] != self._total_num_queries:
                raise ValueError(
                    f"Total number of query tokens mismatch. Expected {self._total_num_queries} from plan(), "
                    f"but got q_nope={q_nope.shape[0]} and q_pe={q_pe.shape[0]} in run()."
                )

@wenscarl wenscarl Dec 20, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This nit comment is valid.

@yzh119 yzh119 Dec 21, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

agreed

@wenscarl wenscarl Dec 21, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure whether to assert q_nope.shape[0] >= self._total_num_queries orq_nope.shape[0] != self._total_num_queries because framework possibly does padding to query. Here is a valid example from EAGLE run where the actual batch 87, but seq_lens is padded to 90.

[INIT METADATA DEBUG] init_forward_metadata called
  forward_mode = 6
  forward_mode.is_decode_or_idle() = False
  forward_mode.is_draft_extend() = True
  forward_mode.is_target_verify() = False
  batch_size = 87
  seq_lens.shape = torch.Size([87])
  spec_info type = EagleDraftInput
  forward_batch.seq_lens = tensor([1275, 1305, 1302, 1331, 1293, 1277, 1288, 1301, 1290, 1313, 1294, 1287,
        1275, 1290, 1272, 1307, 1292, 1284, 1306, 1286, 1272, 1317, 1277, 1281,
        1332, 1287, 1353, 1309, 1329, 1319, 1276, 1274, 1272, 1267, 1285, 1374,
        1365, 1331, 1310, 1297, 1297, 1269, 1291, 1303, 1282, 1365, 1279, 1278,
        1279, 1318, 1336, 1294, 1296, 1299, 1324, 1298, 1301, 1348, 1268, 1327,
        1298, 1312, 1275, 1288, 1281, 1294, 1317, 1287, 1307, 1272, 1295, 1283,
        1293, 1325, 1323, 1295, 1297, 1330, 1297, 1306, 1312, 1286, 1277, 1291,
        1277, 1301, 1279, 1306, 1301,    1], device='cuda:0')
  Total q tokens = 180

For the bug in #2236, it's clearly q_nope.shape[0] < qo_indptr[-1] to cause the IMA.

@bkryu bkryu Dec 21, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given that the q_nope.shape[0] must equal self._total_num_queries assumption is not true, I am thinking that this PR is maybe not necessary and am leaning towards closing the PR without merging.

@wenscarl wenscarl Dec 22, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This check could be handled at the framework level, and closing this PR looks reasonable to me.


if self._backend == "cutlass":
if return_lse:
raise ValueError("return_lse does not support cutlass backend for now.")
Expand Down