trunk/40df7254415ed93bfe83b9d406b44f171ecc5479: Fix flex_attention score_mod with no score gradient (#185991)

FlexAttention builds a joint graph for score_mod so the backward template can compute the gradient of the modified scores with respect to the raw attention scores. When score_mod returns a value that is independent of score, AOTAutograd correctly reports no gradient for the score input. The FlexAttention lowering did not handle that None result: Inductor still expects a score-gradient subgraph output for the dq/dk matmuls, so backward failed during lowering/kernel generation instead of treating the score gradient as zero.

Materialize a zero score gradient when the joint graph returns None for the score input. Constant joint graphs can lower that zero as a scalar/rank-1 Triton value, so the backward template now broadcasts only those low-rank grad_scores values to the score tile before the matmuls. Rank-2 gradients from normal differentiable score_mod paths skip the extra broadcast add.

I considered only adding a template-side fallback, but that would leave the higher-order op contract ambiguous: the joint graph should always provide a score-gradient value to the backward lowering. Materializing the zero in create_fw_bw_graph fixes that contract, while the template rank guard handles the generated representation needed by Triton.

Fixes #145050

Generated by my agent

Benchmark Results:

• CUDA fp16 compiled flex_attention backward, B=1 H=1 S=512 D=64, score_mod returns score * 1.1, 5 warmup iterations and 50 CUDA-event timed iterations.
• Baseline main median: 0.5343 ms; p10/p90: 0.4784/0.6628 ms.
• Patched median: 0.5404 ms; p10/p90: 0.5289/0.5629 ms.
• Result: no clear regression beyond run-to-run noise; baseline timings were bimodal in this environment, while patched timings stayed in the same envelope.

Test Plan:

• Reproduced the issue before the fix with a minimal CUDA repro using torch.compile(flex_attention, dynamic=False), score_mod returning q_idx >= kv_idx, and backward failing with InductorError / joint_subgraph_buffer is None.
• Verified the repro after the fix: q.grad and k.grad are zero, v.grad is nonzero.
• Verified backend="aot_eager" repro after the fix.
• Verified differentiable score_mod=score * score smoke test after the fix.
• python test/inductor/test_flex_attention.py -k test_score_mod_without_score_gradient
• lintrunner -a

Pull Request resolved: #185991
Approved by: https://github.com/drisspg