test: cut unit-test CI wall time - #3601
Conversation
📝 WalkthroughWalkthroughThe PR vectorizes attention, delta-rule, and FP4 test references, and adds a pytest collection hook that prewarms ChangesTest reference vectorization
Mamba checkpointing prewarm hook
Sequence Diagram(s)sequenceDiagram
participant collection_hook as pytest_collection_modifyitems
participant MAX_JOBS_env as MAX_JOBS env
participant build_pool as ThreadPoolExecutor
participant checkpointing_ssu_so as checkpointing_ssu .so
collection_hook->>checkpointing_ssu_so: skip existing shared libraries
collection_hook->>MAX_JOBS_env: override MAX_JOBS for builds
collection_hook->>build_pool: build missing variants in parallel
build_pool->>checkpointing_ssu_so: write built modules
collection_hook->>MAX_JOBS_env: restore original MAX_JOBS
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/bot run |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes kernel compilation and test execution times by removing the batch size B from compile-time cache keys across several GDN decode kernels, allowing a single compiled kernel to be reused across different batch sizes. It also vectorizes reference implementations and page-gathering in tests, shares autotuning buckets, and introduces a parallel JIT pre-compilation fixture for Mamba tests. Feedback recommends dynamically scaling the parallel build workers in the Mamba test fixture based on CPU count to prevent OOM errors on resource-constrained CI runners, and simplifying the aux map in gated_delta_rule_mtp by removing unused default indices.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| prev_max_jobs = os.environ.get("MAX_JOBS") | ||
| os.environ["MAX_JOBS"] = "3" | ||
| try: | ||
| with concurrent.futures.ThreadPoolExecutor(max_workers=12) as pool: | ||
| list(pool.map(lambda spec: spec.build(verbose=False), specs)) |
There was a problem hiding this comment.
On resource-constrained CI runners (such as standard GitHub Actions runners with 2 cores), spawning up to 12 parallel spec builds with MAX_JOBS=3 can result in up to 36 concurrent compiler processes. This can lead to severe CPU thrashing or Out-Of-Memory (OOM) failures. Consider dynamically scaling max_workers and MAX_JOBS based on the available CPU count to ensure optimal parallelization without overloading the system.
cpu_count = os.cpu_count() or 4
max_workers = max(1, cpu_count // 2)
prev_max_jobs = os.environ.get("MAX_JOBS")
os.environ["MAX_JOBS"] = str(max(1, cpu_count // max_workers))
try:
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
list(pool.map(lambda spec: spec.build(verbose=False), specs))There was a problem hiding this comment.
Good catch — 12×3 was tuned on a 32-core box and would oversubscribe a 2-core runner ~18×. Adopted your scaling in next commit.
| "aux": {(B, q.device): (default_indices, default_output)}, | ||
| } | ||
|
|
||
| cache = _compiled_kernels_mtp[cache_key] | ||
| aux_map = cache["aux"] | ||
| aux_key = (B, q.device) | ||
| if aux_key not in aux_map: | ||
| aux_map[aux_key] = ( | ||
| torch.arange(B, dtype=torch.int32, device=q.device), | ||
| torch.empty(B, T, HV, V, device=q.device, dtype=q.dtype), | ||
| ) | ||
| _default_indices_rt, default_output_rt = aux_map[aux_key] |
There was a problem hiding this comment.
In gated_delta_rule_mtp, initial_state_indices is strictly asserted to be not None (line 1923). Therefore, _default_indices_rt is never used at runtime. Storing and allocating torch.arange(B, ...) in the aux map for every new batch size B is unnecessary and wastes GPU memory/allocation overhead. We can simplify the aux map to only store and retrieve default_output.
"aux": {(B, q.device): default_output},
}
cache = _compiled_kernels_mtp[cache_key]
aux_map = cache["aux"]
aux_key = (B, q.device)
if aux_key not in aux_map:
aux_map[aux_key] = torch.empty(B, T, HV, V, device=q.device, dtype=q.dtype)
default_output_rt = aux_map[aux_key]There was a problem hiding this comment.
Good catch — applied. gated_delta_rule_mtp asserts initial_state_indices is not None, so the runtime arange was dead; the aux map now stores only the batch-sized default_output. Fixed in next commit
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
tests/attention/test_trtllm_gen_attention_decode.py (1)
323-324: 💤 Low value
kv_last_page_lenparameter is now unused.The vectorized implementation uses
seq_lensdirectly for token filtering (via the boolean mask at line 351), makingkv_last_page_lenredundant. Consider removing it from the signature to avoid confusion, or add a comment clarifying it's kept for API compatibility.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/attention/test_trtllm_gen_attention_decode.py` around lines 323 - 324, The kv_last_page_len parameter in the function signature is unused because the implementation filters tokens using seq_lens (via the boolean mask), so remove kv_last_page_len: torch.Tensor from the signature and any related references to avoid confusion; update any callers in the test file that pass kv_last_page_len to no longer pass it (or adjust their argument order), and run tests to ensure no other references remain. If you must keep it for API compatibility instead, add a clear comment next to the kv_last_page_len parameter in the signature stating it is intentionally unused and that seq_lens is used for filtering. Use the exact symbol names kv_last_page_len and seq_lens to locate the code to change.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/gemm/test_bmm_fp8.py`:
- Around line 10-16: The test reduced shape coverage by shrinking _TEST_MS and
thus pytest.mark.parametrize("m", _TEST_MS); revert that change so the full
original set of m values is used (restore the original _TEST_MS list used before
this PR) while keeping the shared tuning bucket logic (_TUNING_BUCKETS =
(max(_TEST_MS),)) intact; in practice, replace the current shortened _TEST_MS
with the prior full coverage values and ensure pytest.mark.parametrize("m",
_TEST_MS) continues to drive the test matrix.
- Line 55: The tests call the autotune context manager with
tuning_buckets/round_up unconditionally which installs a shared override even
when auto_tuning is False; update both tests/gemm/test_bmm_fp8.py and
tests/gemm/test_mm_fp4.py so that the autotune(...) invocation only receives
tuning_buckets=_TUNING_BUCKETS and round_up=True when the auto_tuning flag is
True (e.g., conditionally pass those arguments or branch to call
autotune(auto_tuning) vs autotune(auto_tuning, tuning_buckets=_TUNING_BUCKETS,
round_up=True)), ensuring the default bucket mapping is exercised when
auto_tuning is False and referencing the autotune context manager and the
auto_tuning/_TUNING_BUCKETS symbols to locate the change.
In `@tests/mamba/conftest.py`:
- Around line 755-757: In pytest_collection_modifyitems where specs are
prewarmed with ThreadPoolExecutor and spec.build(verbose=False) is invoked, make
the prewarm best-effort by catching and swallowing per-spec exceptions instead
of allowing them to escape and abort collection: replace the direct
pool.map(lambda spec: spec.build(...)) call with a safe wrapper that calls
spec.build(verbose=False) inside a try/except, logs or warns on exception, and
returns/continues so failures fall back to first-touch JIT rather than failing
the whole run; ensure the wrapper references the same spec variable name so it’s
easy to locate.
- Around line 753-756: The test hardcodes ThreadPoolExecutor(max_workers=12)
while setting os.environ["MAX_JOBS"]="3", risking oversubscription; change the
executor to derive max_workers from the MAX_JOBS env (or a sensible default)
and/or cap it against os.cpu_count(), e.g. compute workers =
min(int(os.environ.get("MAX_JOBS", "<default>")), 12, os.cpu_count() or 1) and
pass that to ThreadPoolExecutor so the pool that maps the lambda spec:
spec.build(verbose=False) over specs respects the dynamic MAX_JOBS limit.
---
Nitpick comments:
In `@tests/attention/test_trtllm_gen_attention_decode.py`:
- Around line 323-324: The kv_last_page_len parameter in the function signature
is unused because the implementation filters tokens using seq_lens (via the
boolean mask), so remove kv_last_page_len: torch.Tensor from the signature and
any related references to avoid confusion; update any callers in the test file
that pass kv_last_page_len to no longer pass it (or adjust their argument
order), and run tests to ensure no other references remain. If you must keep it
for API compatibility instead, add a clear comment next to the kv_last_page_len
parameter in the signature stating it is intentionally unused and that seq_lens
is used for filtering. Use the exact symbol names kv_last_page_len and seq_lens
to locate the code to change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7a89b07b-b419-4572-9a5d-de7ff376b9f8
📒 Files selected for processing (10)
flashinfer/gdn_kernels/gdn_decode_bf16_state.pyflashinfer/gdn_kernels/gdn_decode_mtp.pyflashinfer/gdn_kernels/gdn_decode_nontranspose.pyflashinfer/gdn_kernels/gdn_decode_pretranspose.pytests/attention/test_trtllm_gen_attention_decode.pytests/gdn/reference_delta_rule.pytests/gemm/test_bmm_fp8.pytests/gemm/test_mm_fp4.pytests/mamba/conftest.pytests/utils/test_fp4_quantize.py
|
[FAILED] Pipeline #54503417: 8/20 passed |
|
/bot run |
|
/bot stop |
bkryu
changed the title
|
The GitLab CI pipeline #54622394 has been cancelled. |
|
/bot run |
|
[FAILED] Pipeline #54623911: 8/20 passed |
…ze agnostic Extends the BF16-state batch-agnostic fix to the remaining GDN decode kernels (float32 pretranspose, nontranspose, and MTP), so a new running batch size no longer triggers a fresh ~45s cute.compile at inference time. B was a Constexpr in every compile key but unused in all kernel bodies (i_n is decoded from block_idx). B is now a runtime grid extent (q.shape[0]); per-batch tensors are marked dynamic before cute.compile; B (and pool_size for MTP) are dropped from the cache keys; batch-sized aux tensors (h0_indices / cu_seqlens) move to per-(B, device) maps. nontranspose keys on use_small_batch instead of B. Padded/strided pools keep baked strides; T remains constexpr. Mirrors the GDN-kernel portion of Brian Ryu's flashinfer-ai#3601 so that PR can drop its kernel changes and stay test-only.
Aligns the BF16 decode pool handling with the other GDN kernels and with Brian Ryu's flashinfer-ai#3601: a contiguous state pool now marks its slot dimension dynamic (with sentinel -1 pool_size/stride cache keys) so a single cubin serves all pool sizes, while padded/strided pools (vLLM's packed conv+ssm page) keep their real pool_size/stride baked in via the key. Previously B was dropped from the key but pool_size/stride were always retained, so callers with pool_size == batch_size (e.g. the GDN unit tests) still recompiled per batch size. Now compilation is bounded by the tile_v/kernel variant only.
…dynamic batch shapes for vLLM integration) (#3649) ## 📌 Description Addresses part 2 of #3602 , also brings over overlapping effort from @bkryu 's work in #3601 ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance Improvements** * Improved GPU kernel compilation and caching to better reuse compiled cubins across varying batch sizes. * Updated GDN decode/verify (including MTP) kernel launch behavior to derive batch sizing at runtime. * **Robustness Improvements** * Enhanced support for dynamic batch-dependent tensors during compilation and execution. * Improved handling of per-batch auxiliary metadata and correct index/state aliasing for split-pool scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
|
/bot run |
|
/bot run tests/attention |
|
/bot run |
4 participants
📌 Description
Cuts CI wall time for six of the slowest unit-test files. No test cases removed; collection counts, tolerances, and comparison semantics unchanged.
Impacted files
attention/test_trtllm_gen_attention_decode{,_xqa}.pyutils/test_fp4_quantize.pymamba/test_checkpointing_ssu.pyChanges
gdn/reference_delta_rule.py(T×B×heads Python loop → batched einsums),attentionflatten_paged_kv(per-page loop → one gather; also 3 workspace asserts moved GPU-side),utils/test_fp4_quantize.py8x4 scale swizzle (per-element GPU writes → index-math scatter; assertions stayrtol=0, atol=0).mamba/conftest.py(new): pre-builds the 43checkpointing_ssuJIT variants (not in the AOT package) throughJitSpec.build()in a thread pool instead of 43 serial ~60s first-touch compiles. No-op when warm; skipped for small dev runs.🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Bug Fixes
Tests