[ATTN][BUILD] Make the chunk-prefill _b16 tile policy explicit in .conf files

[baodii]

Purpose

Every chunk-prefill config line silently expanded to two translation units: chunk_policy_head<N> and its _b16 variant (same thing with TileShapeQK[1] 32 -> 16, so tiles_per_page stays 1 at page size 16). But the dispatcher only selects _b16 in one case (fmha_xe2.cpp:269):

const bool use_b16_policy = is_paged && (block_size == 16);

So every _b16 kernel with paged=false was dead code: 14 of the 70 in chunk_prefill_default.conf, 60 of the 240 in chunk_prefill_full.conf, each paying full AOT cost across 4 devices. The tile policy was also invisible in the config, so a 35-line .conf producing 70 kernels was not apparent to a reader.

This makes b16 a regular required field:

headsize,paged,causal,local,sink,lse,b16

All seven fields are mandatory and each line maps to exactly one kernel. b16=false builds the standard policy, b16=true builds _b16, and b16=true with paged=false is rejected as undispatchable. The parser requires exactly 7 fields and warns plus skips otherwise, which removes the old implicit expansions that made the kernel set hard to predict from reading a .conf: a bare head size produced 40 kernels from one line, and a 6-field line produced 2.

The same b16 => paged rule is applied to the all keyword, so chunk_prefill_full.conf drops the 60 undispatchable tuples too. This is the config the CI build-wheel job uses, and that job gates run-unit-tests-pvc / run-unit-tests-bmg.

chunk_prefill_default.conf is rewritten fully explicit. Paged decode already treats this dimension as a first-class pagesize field, so this brings chunk prefill to the same footing.

Test Plan

No unit test exists for the CMake parser and this box has no XPU hardware, so the parser was executed directly via a cmake -P harness that stubs configure_file() and prints the generated TU list. Verified the accepted and rejected forms, and diffed the emitted kernel sets old vs new.

Test Result

Config	Before	After
chunk_prefill_default.conf	70	49
chunk_prefill_full.conf	240	180

• All 35 standard-policy TUs are byte-identical to before, so the non-_b16 dispatch surface is unchanged.
• For full.conf, the emitted set was diffed old vs new: 60 removed, 0 added, and all 60 are _b16 with paged=false. Zero undispatchable kernels remain in either preset.
• 7 more _b16 TUs removed from default.conf that no test exercises: head 96 paged (x3), head 192 paged (x2), head 256 paged+LSE, head 512 paged+LSE.
• The 14 retained _b16 TUs in default.conf are all backed by test_flash_attn_varlen_func.py.
• Both presets parse with 0 warnings; files are ASCII-only (file(STRINGS) splits on non-ASCII bytes).

Malformed entries that used to be accepted silently are now rejected:

128                                       -> expected 7 fields, got 1
128,true,true,false,false,false           -> expected 7 fields, got 6
64,true,true,false,false,false,true,extra -> expected 7 fields, got 8
128,true,true,false,false,false,          -> invalid config boolean entry
128,true,,false,false,false,true          -> invalid config boolean entry
192,false,true,false,false,false,true     -> b16=true requires paged=true

Both shipped presets are unaffected by the stricter parser: chunk_prefill_default.conf is already fully explicit and chunk_prefill_full.conf uses the all keyword.

Please review carefully

This is a breaking change to the chunk-prefill config format. Any out-of-tree .conf using a bare head size or the 6-field form will now warn and skip rather than build. Both in-tree presets are fine, but downstream configs need the b16 field added.

The 7 untested _b16 removals from default.conf are a real, if narrow, coverage reduction, and flash_attn_varlen_func swallows "not compiled" errors into a PyTorch fallback, so a missing kernel is a slowdown rather than a failure. Page-16 coverage in the two default configs now disagrees:

• paged_decode_default.conf has pagesize=16 for heads 64, 96, 128, 192, 256
• chunk_prefill_default.conf has b16=true for heads 64, 128, 256, 512

Heads 96 and 192 (Phi / Qwen3.5, DeepSeek MLA) would decode with a real kernel but prefill through the fallback at page size 16. Realigning costs 5 TUs (49 -> 54). Happy to add them back if you prefer symmetry over the smaller build.

The full.conf reduction carries no coverage risk, since the removed kernels are unreachable by construction.

Unrelated finding, for visibility: the UTs exercise 108 distinct chunk-prefill tuples but only 38 exist in chunk_prefill_default.conf. CI never notices because the unit-test jobs install the full-config wheel, while build-wheel-with-default-config is built but never tested. That job is also 100% ccache hits today (measured: 761 cacheable calls, 761 hits, 0 misses), so the default.conf reduction is a cold-build and binary-size win rather than a CI-time win. The full.conf change is the one that shortens the critical path. All of this predates the PR and is otherwise left alone.

(Optional) Documentation Update

• KERNEL_CONFIGURATION.md: documented b16 as a required field with the paged=true rule and the strict field count; corrected the stale default-preset counts (claimed ~31/~17, actual 49/32); updated the full-preset counts to 180.
• csrc/xpu/attn/kernel_configs/README.md: updated counts, replaced the now-wrong "each line expands to two kernels" note, documented what all expands to.
• chunk_prefill_configure.cmake header and both .conf headers describe the seven-field format; chunk_prefill_full.conf had a stale optional-field description that is now corrected.

[Copilot]

🟡 Changes recommended

The updated CMake config parser still accepts malformed field counts (2–5 fields or >7 fields) in ways that can silently expand to unintended kernel sets, and should reject these with a warning to avoid accidental large builds.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR makes the chunk-prefill “_b16 tile policy” axis explicit in chunk-prefill .conf entries, aligning config readability with what actually gets compiled and reducing default-build AOT cost by avoiding undispatchable kernels.

Changes:

• Extend chunk_prefill_configure.cmake to support an optional 7th b16 field, allowing explicit selection of standard vs _b16 policy (and skipping b16=true,paged=false as unreachable).
• Update chunk_prefill_default.conf to be fully explicit (one kernel per line), reducing the default preset from 70 to 49 kernels.
• Refresh documentation to describe the new b16 field and correct kernel-count tables.

File summaries

File	Description
KERNEL_CONFIGURATION.md	Documents the new optional b16 field semantics and updates kernel-count tables.
csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake	Adds parsing/validation and tuple generation for optional b16 policy selection.
csrc/xpu/attn/kernel_configs/README.md	Updates kernel counts and clarifies how b16 affects expansion to translation units.
csrc/xpu/attn/kernel_configs/chunk_prefill_default.conf	Makes all entries explicit and removes unreachable/dead _b16 kernels from the default preset.

Review details

Suppressed comments (1)

csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake:233

• Entries with 2–5 fields (e.g. a missing boolean) currently fall into the “no booleans specified” branch and expand into all 20 combinations for that head size. Given the documented format (either 1 field, 6 fields, or 7 fields), these partial entries should be treated as invalid and skipped with a warning to avoid accidentally compiling a large kernel matrix due to a typo.

            "${b16_policy_${_headsize}}|${_paged}|${_causal}|${_local}|${_sink}|${_lse}"
          )
        endif()
      else()
        # No booleans specified: generate all 20 valid combinations.

• Files reviewed: 4/4 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jikunshang]

full build wo/ cache still take 5 hours in our CI:(
see https://github.com/vllm-project/vllm-xpu-kernels/actions/runs/33851217655/job/100970540100?pr=575