Conversation
📝 WalkthroughWalkthroughAdds caching to the bf16/fp4 runner factories in the CUTE DSL and cuDNN paths. The cuDNN runner now keys cached construction off the effective M-bucket mapper derived from ChangesRunner caching
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![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 introduces caching to the _cute_dsl_bf16_fp4_runner function in flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py using @functools.lru_cache(maxsize=None) to avoid rebuilding the stateless runner class on every API call. The reviewer suggested using @functools.cache instead of @functools.lru_cache(maxsize=None) for cleaner, more modern code and consistency with other parts of the codebase.
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.
_cute_dsl_bf16_fp4_runner defines and instantiates a fresh runner class on every mm_bf16_fp4 call. The runner is stateless (no instance attributes; the closure only captures enable_pdl), so cache the factory on enable_pdl -- same pattern as the cudnn backend's cached graph builders. The per-call class rebuild measures ~5 us (timeit, server x86 core) -- small against the 0.16-0.79 ms kernel at decode shapes, but it is pure avoidable work on the API hot path.
e700562 to
5b57d2d
Compare
bkryu
left a comment
bkryu
left a comment
There was a problem hiding this comment.
Hi @zihaowang-builder, thanks for opening the PR with this change. Left a few minor comments
While we are at it, can you also check whether _cudnn_bf16_fp4_runner() should get the same treatment?
Switch @functools.cache to @functools.lru_cache(maxsize=1024) to match the graph-builder caches in gemm_base.py, and drop the caching rationale paragraph from the docstring.
_cudnn_bf16_fp4_runner rebuilds and re-instantiates the runner class on every mm_bf16_fp4 call (~10-12 us/call measured with timeit on a server x86 core; the cudnn graph builds behind it are already lru_cached). Unlike the cute-dsl factory, this one is not safe to cache on tuning_config: it captures the effective M-bucket mapper from AutoTuner.get().get_effective_map_to_tuning_buckets(), which reflects any thread-local autotune(tuning_buckets=..., round_up=...) override active at call time. A cache keyed on tuning_config would freeze a stale mapper, making the runner's override-shape cache_m bucketing diverge from the bucketing the autotuner uses for cache lookup (10 of 12 probed M values diverged in a repro, e.g. M=1000: autotuner bucket 4096 vs stale cache_m 1024). So split the factory: resolve the effective mapper per call (~0.6 us, required regardless), and lru_cache the runner build keyed on the mapper object -- stable per override scope because _apply_tuning_overrides caches the rewritten TuningConfig. The autotuner profiling-cache key (hash over runner __dict__) keeps its current semantics: stable on the default path, distinct across override scopes. End-to-end mm_bf16_fp4(backend="cudnn") host time at the repo's trace-example shape (M=128, N=2048, K=7168) drops ~165 to ~142 us/call on B200; decode-like M=8 shows the same ~20 us delta. The saving exceeds the ~11 us factory microbench because timeit disables GC while the discarded per-call class objects create reference cycles the real API loop must collect. Validated with tests/gemm/test_mm_bf16_fp4.py (233 passed, 1 skipped) on B200 (sm_100a, cudnn backend 9.24) plus a standalone harness covering bucket consistency inside/outside overrides, cache hits across override scopes, and hash stability.
zihaow211
changed the title
|
@bkryu So I split it: resolve the effective mapper per call (~0.6 µs, needed anyway), and End-to-end it's worth more than the factory microbench suggests: at the trace-example shape (M=128, N=2048, K=7168) the cudnn-backend host time drops ~165 → ~142 µs/call on B200 (same ~20 µs at decode-like M=8) — the extra beyond the ~11 µs factory cost is GC, since the per-call throwaway class objects create reference cycles that Also applied your other two comments (lru_cache(maxsize=1024) + docstring trim). |
|
Hi @yanqinz2 can you review the cuDNN runner changes in this PR? |
|
Hi @yanqinz2 can you help me review the cuDNN runner changes in this PR please? |
3 participants
📌 Description
_cute_dsl_bf16_fp4_runnerand_cudnn_bf16_fp4_runnerdefine andinstantiate a fresh runner class on every
mm_bf16_fp4call — measured~8 µs (cute-dsl) and ~10-12 µs (cudnn) per call of pure avoidable work on
the API hot path (timeit, server x86 core).
cute-dsl: the runner is stateless (the closure only captures
enable_pdl), so the factory is cached directly with@functools.lru_cache(maxsize=1024).cudnn (per review suggestion): this factory is not safe to cache on
tuning_config— it capturesAutoTuner.get().get_effective_map_to_tuning_buckets(...)at build time,which reflects the thread-local
autotune(tuning_buckets=..., round_up=...)override active at that moment. A naive cache would freeze a stale
m_bucket_mapper: measured 10/12 probed M values where the stale runner'sinternal
cache_mbucketing diverges from the autotuner's cache-lookupbucketing (e.g. M=1000: autotuner bucket 4096 vs stale cache_m 1024) —
exactly the hazard
get_effective_map_to_tuning_buckets's docstring warnsabout. Instead the factory is split in two: resolve the effective mapper
per call (~0.6 µs, unavoidable anyway), then cache the runner keyed on the
mapper object — stable per override thanks to
_apply_tuning_overrides'own cache. Verified: bucket-consistent inside/outside overrides, cache hits
within and across override scopes,
hash(runner)(part of the autotunerprofiling-cache key) stable for the default path and distinct across
override scopes.
Net effect: cute-dsl ~8 µs → 0.03 µs; cudnn ~10-12 µs → ~0.9 µs per call
(factory microbench). End-to-end
mm_bf16_fp4(backend="cudnn")host timeat the repo's trace-example shape (M=128, N=2048, K=7168, the
mm_bf16_fp4_cudnn_N2048_K7168_block_size16benchmark definition):~165 µs → ~142 µs per call; decode-like M=8 shows the same ~20 µs delta
(A/B on B200, interleaved 5-round medians). The saving exceeds the factory
microbench because
timeitdisables GC, while the discarded per-call classobjects create reference cycles the real API loop must collect.
Behavior-neutral: the autotuner keys on the runner's class name,
hash(runner), andget_cache_key_extras— all with unchanged values onevery path exercised today.
🔍 Related Issues
N/A
🚀 Pull Request Checklist
✅ 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
Tests have been added or updated as needed.
All tests are passing (
unittest, etc.).Original revision (cute-dsl change only):
pytest tests/gemm/test_mm_bf16_fp4.pygreen on RTX PRO 6000 Blackwell SE (sm_120a), CUDA 13.1.
Current revision (incl. the cudnn runner cache):
pytest tests/gemm/test_mm_bf16_fp4.py— 233 passed, 1 skipped on B200 (sm_100a,cuDNN backend 9.24, cutlass-dsl 4.5), covering both backends. The 1 skip
is the pre-existing parametrize gap "cute-dsl requires out_dtype ==
a.dtype", unrelated to this change.
pytest tests/trace/test_mm_bf16_fp4_reference_correctness.py— 4 passed.Runner-cache semantics validated with a standalone harness: override
bucket consistency (12 M probes × 3 override scopes), instance reuse,
profiling-cache hash stability.
Reviewer Notes
AI-assisted (Claude); I reviewed the change and ran the validation.
Summary by CodeRabbit