Don't schedule prefill unless there is enough kv cache left to fulfil decode stage

[AleksKnezevic]

Ticket

Link to Github Issue

Problem description

When generating high OSL sequences if a prefill is scheduled with low KV cache availability, it will be serviced, then immediately discarded as there is not enough cache available to fulfill decode.

What's changed

Added minimal overhead, defaults to 25%, controllable by TT_XLA_PREFILL_KV_WATERMARK_PERCENT

Checklist

• New/Existing tests provide coverage for changes

[kmabeeTT]

Awesome, looks good - this should help me replace the max_concurrency=32 -> 8 drop I had to make as a workaround in tt-inference-server to get Qwen3-8B running r1_gpqa_diamond eval to take 80 minutes instead of timeout for this effort (tenstorrent/tt-inference-server#4131 (comment))

[mmanzoorTT]

We are mostly using env variable as TTXLA_***

Suggested change

	env_wm = os.environ.get("TT_XLA_PREFILL_KV_WATERMARK_PERCENT")
	env_wm = os.environ.get("TTXLA_PREFILL_KV_WATERMARK_PERCENT")

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 33.82%. Comparing base (9b0c875) to head (11f9734).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5467      +/-   ##
==========================================
- Coverage   33.84%   33.82%   -0.03%     
==========================================
  Files          37       37              
  Lines        4990     4990              
==========================================
- Hits         1689     1688       -1     
- Misses       3301     3302       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.