EP-overlap graph chunking traces FlexAttention with unbacked symbolic batch and sequence dimensions. These are valid runtime sizes, but FlexAttention predicate sites forced trace-time decisions that should either fall back to the general kernel or remain represented as deferred symbolic assertions.
Flex decoding remains an optional optimized implementation. Its predicates use guard_or_false because they are kernel eligibility checks, not user-visible invariants: when ShapeEnv can prove eligibility, decode can be selected; when an unbacked predicate is not provable, the general FlexAttention kernel is used. Failing to prove an optimization must not impose a runtime shape contract on user inputs.
Backward fake propagation preserves the accepted KV-batch-broadcast metadata contract. It reduces grad_key and grad_value back to key/value batch when key/value is provably batch-broadcasted; otherwise it records the non-broadcast invariant with torch._check(Bq == Bkv) and returns unreduced metadata.
Public BlockMask length validation is expressed as two symbolic torch._check predicates: the mask must not be smaller than the query/key lengths, and it must not be larger than those lengths. Concrete mismatches still fail with the existing guidance, while unbacked symbolic lengths become deferred runtime assertions instead of forcing Python to branch on an unbacked expression.
This PR was authored with assistance from an AI assistant.
Test Plan:
python -m pytest test/inductor/test_flex_attention.py -q -s -k "unbacked_flex_decoding_eligibility_falls_back or backward_fake_symbolic_query_key_batch_non_broadcast or block_mask_vs_sequence_lengths"python -m pytest test/inductor/test_flex_attention.py -q -s -k "mask_mod_handles_symint_addition or mask_mod_handles_derived_symint_closure or symbol_closure_in_score_mod"lintrunner -astack-info: PR: #183838, branch: sanketpurandare/stack/12
SOCIAL SHARE CARD GENERATOR