build_backend.py (1)

108-108: Wheel/sdist copies the entire 3rdparty/cccl tree, not just packaged subdirs.

ln("3rdparty/cccl", "cccl") uses shutil.copytree for wheel/sdist (non-symlink) paths, which copies the whole CCCL checkout (CMake files, docs, tests, .github, c/parallel, cudax, python/, etc.) into flashinfer/data/cccl/. Setuptools' package-data globs in pyproject.toml will still filter what ends up inside the built wheel to just cub/cub/**, libcudacxx/include/**, and thrust/thrust/**, so the final .whl isn't bloated — but the intermediate source tree staged during builds grows by hundreds of MB and every _prepare_for_* call does a full rmtree + recopy.

Consider copying only the three header subtrees that are actually packaged (matching the pyproject.toml globs) for the wheel/sdist case, e.g.:

♻️ Suggested refinement

-    ln("3rdparty/cccl", "cccl")
+    if use_symlinks:
+        ln("3rdparty/cccl", "cccl")
+    else:
+        # Only stage the header subtrees that are declared as package-data in pyproject.toml
+        for sub in ("cub/cub", "libcudacxx/include", "thrust/thrust"):
+            ln(f"3rdparty/cccl/{sub}", f"cccl/{sub}")

Note ln's current dst.exists() handling may need a small tweak to create parent directories for nested targets.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@build_backend.py` at line 108, The ln("3rdparty/cccl", "cccl") call copies
the entire CCCL checkout; change ln so that for non-symlink (wheel/sdist) paths
it only copies the three packaged header subtrees (cub/cub, libcudacxx/include,
thrust/thrust) from 3rdparty/cccl into the cccl target instead of copytreeing
the whole repo; update ln's existence handling to ensure parent directories for
nested targets are created before copy/rename operations so dst.exists() checks
work for nested paths; locate and modify the ln function and its call sites to
implement this selective-copy behavior consistent with the pyproject.toml
package-data globs.

flashinfer/jit/cpp_ext.py (1)

98-146: Precedence over CTK CCCL via -I is correct; minor note on extra_include_dirs ordering.

Emitting vendored CCCL as -I before -isystem $cuda_home/include ensures it shadows the CTK-bundled CCCL during header search, which is the stated goal. The #pragma GCC system_header inside CCCL suppresses warnings, so dropping -isystem here is fine. Inline comment with the issue link is helpful.

One subtle ordering note: extra_include_dirs are emitted as -I before the vendored CCCL -I entries (Line 135-142). If a caller ever passes an extra_include_dirs entry that itself contains a (possibly stale/system) cub/, thrust/, or cuda/ header tree, it will now shadow the vendored CCCL. In practice this is unlikely for current callers, but worth keeping in mind — putting CCCL first among the -I flags would be strictly safer:

♻️ Optional: put vendored CCCL ahead of user extras

-    if extra_include_dirs is not None:
-        for extra_dir in extra_include_dirs:
-            common_cflags.append(f"-I{extra_dir.resolve()}")
-    # Vendored CCCL headers use -I (not -isystem) so they take precedence
-    # over the CTK-bundled copy. CCCL headers use `#pragma` system_header
-    # internally to suppress warnings. See https://github.com/NVIDIA/cccl/issues/527
-    for cccl_dir in cccl_includes:
-        common_cflags.append(f"-I{cccl_dir}")
+    # Vendored CCCL headers use -I (not -isystem) so they take precedence
+    # over the CTK-bundled copy. CCCL headers use `#pragma` system_header
+    # internally to suppress warnings. See https://github.com/NVIDIA/cccl/issues/527
+    for cccl_dir in cccl_includes:
+        common_cflags.append(f"-I{cccl_dir}")
+    if extra_include_dirs is not None:
+        for extra_dir in extra_include_dirs:
+            common_cflags.append(f"-I{extra_dir.resolve()}")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/jit/cpp_ext.py` around lines 98 - 146, The current
build_common_cflags emits extra_include_dirs as -I before vendored CCCL -I
entries (cccl_includes), which lets caller-provided includes shadow the vendored
CCCL; to fix, reorder so get_cccl_includes() entries are appended to
common_cflags before iterating extra_include_dirs (or otherwise ensure
cccl_includes are added prior to extra_include_dirs) in build_common_cflags,
referencing the cccl_includes variable and
get_cccl_includes()/build_common_cflags functions to locate the change.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

include/flashinfer/trtllm/fmha/lse.cuh (1)

29-38: Add a brief inline justification for cub::DeviceTransform and the dropped launch_with_pdl.

The kernel used to be hand-rolled and its sole caller (fmhaKernels.cuh:285-291) still passes params.enable_pdl as the 4th argument, which is now silently discarded. The PR description notes that cub::DeviceTransform already uses PDL/vectorized loads/prefetch/TMA internally — that rationale belongs in the header next to the functor so future readers understand (a) why the parameter is intentionally unused rather than wired through, and (b) why the custom launcher was removed. Per coding guidelines, performance-critical hot paths should document special algorithmic choices and alternatives considered.

💡 Proposed comment

+// Use cub::DeviceTransform instead of a hand-rolled kernel: it auto-tunes
+// occupancy and internally applies PDL, vectorized loads, prefetch, and TMA
+// where available. Because CUB manages PDL itself, `launch_with_pdl` from
+// the caller is intentionally ignored (kept for API stability).
 inline cudaError_t ComputeLSEFromMD(float2* md, float* lse, int n, bool /*launch_with_pdl*/,
                                     cudaStream_t stream) {
   return cub::DeviceTransform::Transform(md, lse, n, MDToLSE{}, stream);
 }

As per coding guidelines: "For performance-critical hot paths, leave comments with justification for special algorithmic choices and mention alternative approaches considered".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@include/flashinfer/trtllm/fmha/lse.cuh` around lines 29 - 38, Add a concise
inline comment above the MDToLSE functor and ComputeLSEFromMD function
explaining why cub::DeviceTransform is used instead of the previous hand-rolled
kernel and why the boolean parameter launch_with_pdl (and callers passing
params.enable_pdl) is intentionally ignored: state that cub::DeviceTransform
provides PDL/vectorized loads, prefetch/TMA and optimized device transforms
internally so the custom launcher was removed for maintainability and equivalent
performance; mention that callers still pass enable_pdl for API compatibility
and that it is intentionally unused here. Ensure the comment references MDToLSE
and ComputeLSEFromMD by name and briefly notes that alternatives considered
included keeping the hand-rolled kernel or forwarding launch_with_pdl to a
custom launcher.

flashinfer/jit/cpp_ext.py (1)

98-142: LGTM on -I placement and rationale; consider surfacing a clearer error when CCCL paths are absent.

The ordering (extra_include_dirs -I → cccl -I → system -isystem) correctly gives vendored CCCL precedence over $cuda_home/include, and the inline comment pointing at CCCL #527 explains the -I vs -isystem choice well.

One minor robustness nit: Path.resolve() (non-strict by default) will happily produce a canonicalized path for directories that don't exist. In a source tree where 3rdparty/cccl hasn't been initialized (e.g., git clone without --recurse-submodules and not an installed wheel), JIT compilation will fail inside nvcc with a header-not-found error instead of a clear message that the submodule is missing. Consider validating at least one expected CCCL header path exists in get_cccl_includes() and raising a descriptive error pointing the user to git submodule update --init 3rdparty/cccl.

💡 Optional hardening

 def get_cccl_includes() -> List:
     """Get vendored CCCL include directories (added with -I for CTK override precedence)."""
-    return [p.resolve() for p in jit_env.CCCL_INCLUDE_DIRS]
+    resolved = [p.resolve() for p in jit_env.CCCL_INCLUDE_DIRS]
+    missing = [str(p) for p in resolved if not p.is_dir()]
+    if missing:
+        raise RuntimeError(
+            "Vendored CCCL headers not found at: "
+            + ", ".join(missing)
+            + ". Initialize the submodule with "
+            "`git submodule update --init --recursive 3rdparty/cccl`."
+        )
+    return resolved

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/jit/cpp_ext.py` around lines 98 - 142, get_cccl_includes currently
returns resolved paths even if the CCCL files are missing, which leads to a
confusing nvcc header-not-found error later; modify get_cccl_includes to check
that at least one expected CCCL header (e.g., a canonical file under
jit_env.CCCL_INCLUDE_DIRS such as "cccl.h" or another known header) exists on
disk after resolving, and if none exist raise a clear RuntimeError explaining
the CCCL submodule is missing and instructing the user to run "git submodule
update --init 3rdparty/cccl" (or to install the wheel) so the JIT compile fails
fast with a helpful message; reference the get_cccl_includes function and
jit_env.CCCL_INCLUDE_DIRS when adding the existence checks and error raise.