[Perf][Qwen4Exp] Overlap the PLE gather of the next prefill chunk

docs/design/architecture/runtime_reference.md

@@ -38,7 +38,7 @@ identify initialization/hash code; they do not evaluate its semantics.
 | `vllm.config.gdn_state.GdnStateConfig` | kernel_config.gdn.state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L39); graph_options: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L52) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L23) |
 | `vllm.config.gdn_state.GdnStateTraceConfig` | observability_config.gdn_state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L99); compile_ignored_aliases: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L142) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L59) |
 | `vllm.config.kernel.IrOpPriorityConfig` | kernel_config.ir_op_priority | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L84) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L67) |
-| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1044) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) |
+| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1048) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) |
 | `vllm.config.kernel.Sm70AwqConfig` | kernel_config.sm70_awq | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L338); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L317) |
 | `vllm.config.kernel.Sm70Fp8Config` | kernel_config.sm70_fp8 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L433); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L369) |
 | `vllm.config.kernel.Sm70GgufConfig` | kernel_config.sm70_gguf |  | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L498) |

tests/compile/test_sm70_decode_graph.py

@@ -301,6 +301,8 @@ def test_qwen38_hybrid_ple_skips_decode_offload_request(monkeypatch) -> None:
     connector = SimpleNamespace(
         _ple_policy=ple_policy(),
         _launch=lambda num_reqs, num_tokens: launches.append((num_reqs, num_tokens)),
+        early_prefill=False,
+        _host_result_regions={},
     )
 
     PleOffloadConnector.prepare_forward(
@@ -376,7 +378,10 @@ def test_hybrid_ple_admits_target_decode_widths(monkeypatch, tokens):
     monkeypatch.setenv("VLLM_SM70_QWEN38_HYBRID_PLE", "1")
     launches = []
     connector = SimpleNamespace(
-        _ple_policy=ple_policy(), _launch=lambda *args: launches.append(args)
+        _ple_policy=ple_policy(),
+        _launch=lambda *args: launches.append(args),
+        early_prefill=False,
+        _host_result_regions={},
     )
     PleOffloadConnector.prepare_forward(connector, 1, tokens, False, True)
     assert launches == []
@@ -419,7 +424,7 @@ def get_capability(device_id):
     monkeypatch.setattr(cg, "graph_capture", lambda **kwargs: nullcontext())
     monkeypatch.setattr(cg, "is_global_first_rank", lambda: False)
     monkeypatch.setattr(torch.cuda, "CUDAGraph", object)
-    monkeypatch.setattr(torch.cuda, "graph", lambda *args: nullcontext())
+    monkeypatch.setattr(torch.cuda, "graph", lambda *args, **kwargs: nullcontext())
     monkeypatch.setattr(
         cg,
         "get_offloader",
@@ -430,6 +435,7 @@ def get_capability(device_id):
     manager = object.__new__(cg.CudaGraphManager)
     manager.device = torch.device("cpu")
     manager.pool = None
+    manager.capture_context = SimpleNamespace(stream=None)
     manager.graphs = {}
     manager._capture_descs = {
         mode: [

tests/kernels/ple/test_host_result.py

@@ -101,7 +101,6 @@ def setup_cross_process_offload(self, output, sem):
     monkeypatch.setattr(
         module.torch, "empty", lambda *a, **kw: allocate(*a, **{**kw, "device": "cpu"})
     )
-    monkeypatch.setattr(module.envs, "VLLM_SM70_QWEN38_HYBRID_PLE", False)
     monkeypatch.setattr(
         module.HostResultRegion,
         "create",
@@ -120,6 +119,7 @@ def setup_cross_process_offload(self, output, sem):
             ple_result_transport="auto", ple_result_transports={}
         ),
         speculative_config=None if method is None else SimpleNamespace(method=method),
+        offload_config=SimpleNamespace(ple=SimpleNamespace(hybrid=False)),
     )
     connector = module.PleOffloadConnector.__new__(module.PleOffloadConnector)
     connector.device = torch.device("cpu")

tests/models/qwen4_exp/test_gguf_pinned_policy.py

@@ -84,7 +84,7 @@ def remote(*args):
         forward_impl=local,
         wait_offloaded_output=remote,
     )
-    monkeypatch.setattr(module.envs, "VLLM_SM70_QWEN38_HYBRID_PLE", False)
+    monkeypatch.delenv("VLLM_SM70_QWEN38_HYBRID_PLE", raising=False)
     monkeypatch.setattr(module, "use_sm70_decode_graph_semantics", lambda: decode)
     result = module.PleOffloadLayer.forward(layer, None, SimpleNamespace(shape=(512,)))
     assert result is marker
@@ -104,9 +104,11 @@ def test_connector_only_skips_requests_for_owned_local_decode(
     calls = []
     connector = SimpleNamespace(
         _all_pinned_decode=pinned,
+        _ple_policy=SimpleNamespace(hybrid=False),
         _launch=lambda *a: calls.append(a),
+        early_prefill=False,
+        _host_result_regions={},
     )
-    monkeypatch.setattr(module.envs, "VLLM_SM70_QWEN38_HYBRID_PLE", False)
     module.PleOffloadConnector.prepare_forward(
         connector, 4, 20, dummy_run=False, use_local_model=local_model
     )

tests/models/qwen4_exp/test_ple.py

@@ -734,7 +734,7 @@ def test_ngram_embedding_disk_offload_allocates_only_meta_table(
     monkeypatch.setattr(
         parameter_module, "get_tensor_model_parallel_world_size", lambda: 1
     )
-    monkeypatch.setattr(ple_module.envs, "VLLM_PLE_DISK_OFFLOAD", True)
+    monkeypatch.setenv("VLLM_PLE_DISK_OFFLOAD", "1")
     monkeypatch.setattr(ple_module.envs, "VLLM_PLE_DISK_OFFLOAD_NUM_THREADS", 0)
     monkeypatch.setattr(ple_module, "is_offload_process", lambda: True)
     config = SimpleNamespace(
@@ -1045,7 +1045,7 @@ def test_ngram_gpu_offload_retains_only_fp8_global_scale(monkeypatch) -> None:
     nn.Module.__init__(module)
     module._offload_model_dtype = torch.float16
     weight_scale = torch.tensor([0.25], dtype=torch.bfloat16)
-    monkeypatch.setattr(ple_module.envs, "VLLM_PLE_CPU_OFFLOAD", True)
+    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
     monkeypatch.setattr(ple_module, "is_offload_process", lambda: False)
     monkeypatch.setattr(
         torch.accelerator,
@@ -1294,6 +1294,7 @@ def test_ngram_cpu_offload_padding_does_not_overwrite_real_tokens(
     module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
     nn.Module.__init__(module)
     module._packed_gguf = False
+    module.layer_name = "model.layers.0.ple"
     module.embedding_dim = 1
     module.head_dim = 1
     module.ngram_size = 2
@@ -1402,6 +1403,7 @@ def test_ngram_fp8_cpu_offload_preserves_quantized_output(
     module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
     nn.Module.__init__(module)
     module._packed_gguf = False
+    module.layer_name = "model.layers.0.ple"
     module.embedding_dim = 2
     module.head_dim = 2
     module.ngram_size = 2

tests/v1/worker/test_early_model_inputs.py

@@ -31,8 +31,9 @@ def prepare_inputs(batch, states):
         calls.append("inputs")
         return inputs
 
-    def prepare_forward(reqs, tokens, dummy_run):
-        assert (reqs, tokens, dummy_run) == (1, 5, False)
+    def prepare_forward(reqs, tokens, dummy_run, cpu_inputs, eager_forward):
+        assert (reqs, tokens, dummy_run, cpu_inputs) == (1, 5, False, None)
+        assert eager_forward == (mode == CUDAGraphMode.NONE)
         calls.append("ple")
 
     runner = SimpleNamespace(
@@ -41,7 +42,9 @@ def prepare_forward(reqs, tokens, dummy_run):
         ),
         req_states=None,
         _ple_offload_connector=(
-            SimpleNamespace(prepare_forward=prepare_forward) if ple else None
+            SimpleNamespace(prepare_forward=prepare_forward, early_prefill=False)
+            if ple
+            else None
         ),
     )
     batch = SimpleNamespace(num_reqs=1, num_tokens_after_padding=5)

tests/v1/worker/test_ple_offload_worker.py

@@ -397,7 +397,9 @@ def test_ple_offload_uses_event_per_inflight_mrv2_batch() -> None:
     connector = PleOffloadConnector.__new__(PleOffloadConnector)
     connector.tp_rank = 0
     connector.dp_rank = 0
+    connector.fast_lane = None
     connector._uses_cuda_inputs = True
+    connector._next_seq = 0
     connector._request_queue = queue.Queue(maxsize=2)
     connector._d2h_event_pool = queue.Queue(maxsize=2)
     first_event = Mock()
@@ -604,7 +606,7 @@ def test_ple_offload_requires_ple_layers(
     worker.model_config = SimpleNamespace(  # type: ignore[assignment]
         hf_text_config=SimpleNamespace(ple_layer_ids=ple_layer_ids)
     )
-    monkeypatch.setattr(envs, "VLLM_PLE_CPU_OFFLOAD", True)
+    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
 
     assert worker._has_ple_layers() is expected
 
@@ -861,6 +863,7 @@ class FakeStream:
             dtype=torch.float32,
             hf_text_config=SimpleNamespace(ple_embed_dim=2),
         ),
+        kernel_config=SimpleNamespace(ple_fast_lane=False),
     )
     runner._layers = {
         "ple": SimpleNamespace(