feat(gemma4): generate the hd512 global prefill kernel and serve it behind a knob - #1052
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d71c9a3cf
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| HEADS = 32 | ||
| GROUPS = 8 |
There was a problem hiding this comment.
When PEGAINFER_GLOBAL_ATTN=tilelang is used with the supported 12B checkpoint, the engine accepts the knob because startup checks only whether any generated body exists, but this kernel is compiled exclusively for 32 query heads and 4 global KV heads; the published 12B geometry is 16/1, so the launcher's head-count guard returns cudaErrorInvalidValue on the first global prefill and the request fails. Generate and dispatch every supported geometry, or reject incompatible checkpoints during engine loading rather than accepting a configuration that cannot serve a prompt.
Useful? React with 👍 / 👎.
| fn tilelang_arch(sm_targets: &[String]) -> Option<String> { | ||
| let max_sm = sm_targets | ||
| .iter() | ||
| .filter_map(|sm| sm_numeric_prefix(sm)) |
There was a problem hiding this comment.
When PEGAINFER_CUDA_SM lists multiple targets, or the build host exposes mixed GPU generations, selecting only the maximum SM causes tilelang_gencode to emit SASS solely for that architecture. The crate-level gemma4_tilelang cfg still reports the kernel as built, so selecting one of the lower devices with PEGAINFER_GLOBAL_ATTN=tilelang passes startup validation and later fails at the function attribute or launch with no compatible kernel image. Generate an image for each supported requested SM or validate the selected runtime device before accepting the knob.
AGENTS.md reference: AGENTS.md:L39-L39
Useful? React with 👍 / 👎.
| lines = [f"CU_PATH={cu_path}"] | ||
| lines.append(f"TILELANG_TEMPLATE_PATH={template_include}") | ||
| lines.append(f"CUTLASS_INCLUDE_DIR={cutlass_include}") |
There was a problem hiding this comment.
When --vendor-includes output is copied to a TileLang-free build host—the purpose of the pre-generated tier—these manifest entries retain the producer's absolute paths, or paths relative to the producer's working directory. parse_tilelang_manifest resolves them directly rather than relative to manifest.txt, so moving the supposedly self-contained directory makes its CUDA file and vendored include roots appear missing and aborts the build. Store paths relative to the artifact directory and resolve them against the configured pre-generated directory.
Useful? React with 👍 / 👎.
b6b1554 to
51d103d
Compare
|
@codex review again plx |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 51d103d4c8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for line in text.lines() { | ||
| if let Some(value) = line.strip_prefix("CU_PATH=") { | ||
| cu_files.push(PathBuf::from(value.trim())); | ||
| cu_files.push(base.join(value.trim())); |
There was a problem hiding this comment.
When using the documented K3 pre-generated flow (generate.py --out-dir pregen --vendor-includes followed by PEGAINFER_K3_TILELANG_PREGEN=pregen), the unchanged K3 generator writes entries such as CU_PATH=pregen/... and TILELANG_TEMPLATE_PATH=pregen/include/...; unconditionally joining those entries to the manifest directory turns them into pregen/pregen/..., so parse_tilelang_manifest reports missing CUDA files and the offline K3 build fails. Make K3 emit manifest-relative paths as the Gemma generator now does, or retain compatibility with its producer-relative entries.
AGENTS.md reference: AGENTS.md:L163-L163
Useful? React with 👍 / 👎.
| for knob in ["PYTHON", "PREGEN", "JOBS"] { | ||
| println!("cargo:rerun-if-env-changed=PEGAINFER_{label}_TILELANG_{knob}"); | ||
| } |
There was a problem hiding this comment.
On an incremental Gemma4 or K3 build that previously selected the stub tier, setting PEGAINFER_TILELANG_PYTHON does not cause Cargo to rerun this build script because only the family-specific PEGAINFER_<label>_TILELANG_PYTHON variables are registered here, even though find_tilelang_python also reads the shared variable. The stale stub therefore remains and PEGAINFER_GLOBAL_ATTN=tilelang continues to be refused until the user cleans the build or changes a tracked input; add a rerun-if-env-changed entry for the shared selector.
Useful? React with 👍 / 👎.
b179c61 to
a6d2e51
Compare
|
@codex review again |
There was a problem hiding this comment.
💡 Codex Review
pegainfer/pegainfer-kernels/build.rs
Lines 1700 to 1703 in a6d2e51
Skip TMA generation on pre-Hopper targets
When Gemma 4 is built for a supported pre-Hopper device such as SM89 and the discovered Python happens to have TileLang installed, this returns sm_89 and generate_tilelang_artifacts invokes the Gemma generator instead of selecting the stub tier. That generator requires TMA descriptors and explicitly errors when lowering produces none, while TMA is unavailable before SM90, so merely installing TileLang makes the otherwise-supported Gemma build fail. Return None for Gemma targets below SM90 or fall back to the stub when this TMA-only generation is inapplicable.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| { | ||
| let refusal = crate::engine::tilelang_geometry_refusal(&config) | ||
| .expect_err("a checkpoint the kernel was not compiled for must be refused"); | ||
| eprintln!("geometry {geometry:?} is not this checkpoint's: {refusal}"); | ||
| return; |
There was a problem hiding this comment.
When the maintained gate suite is run with its usual 12B checkpoint (16 query heads over 1 global KV head), the generated 32/4 geometry mismatches and this branch returns success after testing only startup refusal. Consequently the test named as the production-path comparison never executes the replacement kernel, leaving the output-changing path with only synthetic kernel coverage; run this gate with a compatible checkpoint or fail/skip it explicitly rather than counting the refusal as the model-level parity evidence.
AGENTS.md reference: AGENTS.md:L122-L124
Useful? React with 👍 / 👎.
a6d2e51 to
c9bac8c
Compare
|
@codex review |
c9bac8c to
98883f2
Compare
…ehind a knob Signed-off-by: Feathbow <feathbow@gmail.com>
98883f2 to
8a5d45c
Compare
1 participant
Description
Closes #1051
The shape was never the problem; the budget was. FlashInfer's FA2 path serves head_dim 512 by time-sharing K and V in one shared-memory buffer with a 64-key tile, a layout its own comment ties to the 99 KB SKUs. On a GH200 that leaves the kernel at 19–28% of the measured roofline where FlashAttention 4 holds 44–49%, and at 31B the ten global layers' prefill is the largest account of a long prompt.
A generated kernel, built where the line lives.
pegainfer-gemma4/kernels/carries a TileLang definition of the varlen paged prefill at head_dim 512 and a generator that emits the CUDA, a launcher table and a host stub at build time; the kernels crate's AOT section now takes a family table, so this family sits beside the routed lane's rather than duplicating its build logic. A build without TileLang carries a stub that refuses to serve, and says so.The global pool is paged at the key block. The global family's page size moves to 64 so one key tile is one TMA copy; the sliding pool is untouched in this change.
One knob, three lines it can fall on.
PEGAINFER_GLOBAL_ATTN=tilelangroutes the global family's prompt rows through the generated kernel; unset oroffkeeps the incumbent and byte-identical serving. The walk inside the kernel is bounded by the CTA count the host computed for the plan, so a CTA past the plan does nothing rather than reading past it.Measured, not assumed. A paired harness (
hd512_prefill_bench) times both kernels on one pool and one plan; the serving oracle holds the generated kernel against the incumbent on the checkpoint.Test Env
Verification
Bound to the tip of this range; every run on the box above.
cargo fmt --check,clippy -D warningson the kernels and gemma4 crates including tests and examples, lib tests green.