Interleave chunked prefills with single decoding steps (#558)

# Description

* Deprioritize chunked prefill when previous step was already a prefill
and requests are decoding
* Evaluate blocks condition at the time of first chunked prefill
* Evaluate remaining conditions at the time of last chunked prefill

* Tests will be implemented other PRs:
* interleaving logic correctness: e.g. that after step X which was a
prefill, step X+1 is a decode
* ✅ DONE constraint correctness: verify that we cannot schedule first
chunked prefill when there isn't enough blocks

---------

Signed-off-by: Sophie du Couédic <[email protected]>
Signed-off-by: Travis Johnson <[email protected]>
Co-authored-by: Travis Johnson <[email protected]>

Author: sducouedic

vllm_spyre/envs.py

@@ -13,6 +13,7 @@
     VLLM_SPYRE_PERF_METRIC_LOGGING_DIR: str = "/tmp"
     VLLM_SPYRE_OVERRIDE_SIGNALS_HANDLER: bool = False
     VLLM_SPYRE_USE_CHUNKED_PREFILL: bool = False
+    VLLM_SPYRE_CP_INTERLEAVE_STEPS: bool = True
     # Prompt logprobs are behind a flag because they're only supported for
     # static batching and require passing back the hidden states for the full
     # prefill on every request. This could incur a heavy performance penalty in
@@ -172,6 +173,12 @@ def _backend_backwards_compat() -> str:
     # single prefill is used.
     "VLLM_SPYRE_USE_CHUNKED_PREFILL":
     lambda: bool(int(os.getenv("VLLM_SPYRE_USE_CHUNKED_PREFILL", "0"))),
+
+    # Feature Flag
+    # Works only with chunked prefill enabled. If set, prefill steps are
+    # interleaved with a decode step
+    "VLLM_SPYRE_CP_INTERLEAVE_STEPS":
+    lambda: bool(int(os.getenv("VLLM_SPYRE_CP_INTERLEAVE_STEPS", "1"))),
 }
 # --8<-- [end:env-vars-definition]