Skip to content

[CK_TILE] batch prefill codegen hard-codes kSkipMinSeqlenQ to false while the dispatch key records the requested value #12097

Description

@JinHe-102

Summary

dispatcher/codegen/fmha/codegen.py emits a literal false for
kSkipMinSeqlenQ_ in both batch prefill traits lists, ignoring
sig["skip_min_seqlen_q"]. Every neighbouring flag in those same lists
is emitted as _bool_cpp(sig[...]).

The generated dispatch key, however, does record the requested value.
So codegen exits 0 and registers a kernel under a key claiming the flag
is on, compiled with it off.

Same defect class as #12093 / PR #12092, about 40 lines away in the same
file. Currently latent — see "Severity" below.

Where it happens

traits list template opens offending literal
ck_tile::TileFmhaBatchPrefillTraits<...> line 773 line 784 (slot 11)
fmha_fwd_batch_prefill_traits_<...> line 818 line 835

The key, at line 1230, is correct:

key.signature.skip_min_seqlen_q = {str(sig["skip_min_seqlen_q"]).lower()};

By contrast, _bool_cpp(sig["skip_min_seqlen_q"]) is used at lines 233,
284, 360 and 408 for the fwd and pagedkv families. Batch prefill is
the odd one out.

Slot positions confirmed against
include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp:65-77 and
example/ck_tile/01_fmha/fmha_fwd.hpp:1537-1560.

One caveat for whoever patches this: each list contains a second
false literal that is entirely legitimate and should be left alone —
kHasBiasGrad at line 779, and kUseTrLoad at line 834.

Reproduction

Codegen run twice on a clean 9bb77c148e checkout with only
skip_min_seqlen_q changed, reading both the kernel header and the
dispatcher_wrapper_*.hpp that carries the key. fwd is swept alongside
as a control, since it should track the flag:

family = batch_prefill   (slot 11 = kSkipMinSeqlenQ)
  skip_min_seqlen_q=False -> traits slot 11 = false | key ... = false;
  skip_min_seqlen_q=True  -> traits slot 11 = false | key ... = true;
  => slot 11 IGNORES the signature (both false)

family = fwd             (slot 11 = kSkipMinSeqlenQ)
  skip_min_seqlen_q=False -> traits slot 11 = false | key ... = false;
  skip_min_seqlen_q=True  -> traits slot 11 = true  | key ... = true;
  => slot 11 tracks the signature correctly

The second batch prefill row is the bug. The two runs also produce
different name hashes (5c71ab0638d5 vs 55c564e7a881), so the
dispatcher ends up storing two distinct kernels whose behaviour is
identical.

Severity: latent, not currently reachable

skip_min_seqlen_q is not part of the batch prefill emission surface. It
is absent from the config objects expand_sweep produces for the family,
on both gfx1100 and gfx950, across every
tile_engine/ops/fmha/configs/*.json. The literal false therefore
agrees with everything the shipped sweeps emit today.

It becomes reachable via a hand-written --config-json, which is how we
found it, or if the flag is ever added to the batch prefill sweep.

We are filing it because it is the same shape as #12093: a traits slot
that does not carry its signature value. That one was load-bearing,
because the omission shifted every later argument. This one is silent,
because the literal happens to have the right type and a value that
matches current emissions.

Suggested fix

Either emit _bool_cpp(sig["skip_min_seqlen_q"]) in both lists, or
reject the flag for this family so the key cannot claim something the
kernel does not honour.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions