Repository navigation
[FLYDSL] Support paged mqa logits fp4 varqlen kernel - #4230
Merged
Merged
Conversation
Default parallel_unit_num to None and derive it in compute_varctx_schedule as batch*next_n*ceil(max_seq_len/block_k) when unset. The value comes purely from static shape metadata (no context_lens read), so it is cudagraph-safe and is a multiple of next_n and >= batch*next_n by construction. Callers can now omit the knob, and it fixes next_n values (e.g. next_n=10) where the old default of 512 was not a multiple of next_n. The op-test passes None through to exercise the auto path. Co-authored-by: Cursor <cursoragent@cursor.com>
…ill kernel Add flydsl_pa_mqa_logits_fp4_varqlen + _build_varqlen_mtp_windows: a cu_seq_q (per-batch qlen prefix-sum) adapter over the ragged-prefill kernel, so each batch can have a different query length. Row n of batch b scores KV [0, ctx_b - qlen_b + n] (MTP tail-causal, same semantics as the uniform-next_n decode kernel). The builder uses only static shape metadata (no context_lens value read), so it is cudagraph-safe; no kernel or ABI change is needed. Export the new entry from aiter.ops.flydsl. Merge the standalone varqlen op-test into test_flydsl_pa_mqa_logits_fp4_prefill.py (shared fp4/preshuffle/reference helpers; the builder is unit-checked against an independent window construction). The varqlen cases also report TFLOPS / GB/s (kernel-only, useful FLOPs/bytes), gated behind --bench like the prefill perf. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds variable query-length (per-batch MTP) support to the FlyDSL FP4 paged MQA logits path, including an exported flydsl_pa_mqa_logits_fp4_varqlen wrapper and updated scheduling defaults to allow cudagraph-safe auto-derived CTA grids.
Changes:
- Make
parallel_unit_numoptional (None= auto-derived) for the decode/varctx scheduling path and update the corresponding op-test defaults/docs. - Introduce
flydsl_pa_mqa_logits_fp4_varqlen(cu-seq prefix-sum adapter) plus a helper to build MTP tail-causal windows. - Extend the prefill op-test to validate varqlen correctness and optionally report kernel-only TFLOPS / GB/s.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
op_tests/test_flydsl_pa_mqa_logits_fp4.py |
Switches schedule CTA target to optional/auto-derived and updates CLI help. |
op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py |
Adds varqlen test coverage and perf reporting for the new varqlen path. |
aiter/ops/flydsl/kernels/pa_mqa_logits_fp4.py |
Allows compute_varctx_schedule + decode wrapper to accept parallel_unit_num=None and documents the auto-derivation. |
aiter/ops/flydsl/kernels/pa_mqa_logits_fp4_prefill.py |
Adds _build_varqlen_mtp_windows and flydsl_pa_mqa_logits_fp4_varqlen wrapper. |
aiter/ops/flydsl/__init__.py |
Exports the new public FlyDSL symbol flydsl_pa_mqa_logits_fp4_varqlen. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Fuse the cu_seq_q -> (row_to_batch, local_starts, local_ends) construction into a single Triton kernel (_varqlen_windows_kernel), replacing the ~8 small torch ops (arange/searchsorted/gathers/arith/clamp). Still cudagraph-safe: the grid comes from total_q (q.shape[0]) with no context_lens value read. Rename the builder to compute_varqlen_windows to match compute_prefill_schedule / compute_varctx_schedule, and export it. Decouple window-building from the launch so the ragged-row metadata is built once and reused across many calls: cu_seq_q/context_lens are now keyword-only (one-shot build), with optional windows=(row_to_batch, local_starts, local_ends) and cta_info/n_ctas passthrough for repeated launches with a stable layout. The op-test uses the new signature and adds a windows= reuse check (torch.equal against the cu_seq_q path). Co-authored-by: Cursor <cursoragent@cursor.com>
Sort the flydsl guarded import block and its __all__ additions, and switch the varqlen entry point plus the test's local import to the forms ruff expects. Only diff-context findings are addressed; the files' pre-existing warnings are left untouched. Co-authored-by: Cursor <cursoragent@cursor.com>
valarLip
approved these changes
Jul 28, 2026
valarLip
added a commit
that referenced
this pull request
Jul 28, 2026
Conflicts came from #4230 (FLYDSL paged mqa logits fp4), which touched five files this branch had also reformatted. Resolved by taking main's content verbatim for every conflicted file and re-running the lint pass over it, rather than hand-merging formatting against functional changes: git checkout --theirs <files> && ruff check --fix && black Only mechanical rules applied -- UP045 (5), I001 (2), RUF022 (1), all safe fixes, no --unsafe-fixes needed. `git diff origin/main` on those five files now shows lint changes and nothing else, which is the proof that main's functional content survived intact.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.