[Perf][DSv4.1] Run the Engram WKV ahead for GPU-resident tables too

tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py

@@ -775,15 +775,17 @@ def test_fp32_index_key_store(ratio):
 # --------------------------------------------------------------- Engram
 
 
-@pytest.mark.parametrize("masked", [False, True])
-def test_engram_wkv_prefetch_matches_inline(masked):
-    """The Engram WKV started on a side stream right after its (slow) host
-    lookup, replayed from a CUDA graph, gives the stock forward's output
+@pytest.mark.parametrize(
+    "cpu_offload,masked", [(True, False), (True, True), (False, False)]
+)
+def test_engram_wkv_prefetch_matches_inline(cpu_offload, masked):
+    """The Engram WKV started on a side stream right after its (slow) inline or
+    host lookup, replayed from a CUDA graph, gives the stock forward's output
     bitwise."""
     from tests.kernels.test_engram import _make_embedding
 
     hc_mult, dim, max_tokens = 4, 256, 16
-    table = _make_embedding(cpu_offload=True)
+    table = _make_embedding(cpu_offload=cpu_offload)
     lookup = table.lookup
 
     def slow_lookup(indices, out, background=False):

vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md

@@ -110,7 +110,7 @@ forward runs the stock layers, which the executor does not modify.
 | Compressed layers sharing an index source reuse one sparse-index table; the first of them builds the other KV-cache groups' tables on a side stream | all executor steps |
 | Sparse-MLA decode on a CuTe DSL kernel, a thread block cluster per token with one CTA per 128 columns (`attention/sparse_mla_cute.py`) | the target's causal decode while every token's cluster is resident at once (at width 640, up to 8 requests with DSpark on a GB300) |
 | Otherwise the TRTLLM-gen sparse-MLA decode, launched on a KV-counter buffer zeroed once | all other executor steps |
-| Engram WKV chain on a side stream right after each layer's lookup | all executor steps, with host-offloaded TP-only tables and a CuTe DSL MXFP8 `wkv` |
+| Engram WKV chain on a side stream right after each layer's lookup | all executor steps, with TP-only tables (GPU-resident or host-offloaded) and a CuTe DSL MXFP8 `wkv` |
 | Attention metadata builds recorded at CUDA graph capture, one per batch shape the FULL decode graphs serve, and replayed | FULL-graph uniform decode |
 | LM head on FlashInfer's CuTe DSL BF16 GEMM | 1-16 rows |
 | Logits gathered with the executor's Lamport all-gather | up to 4 MiB a rank (at TP4, 32 FP32 or 64 BF16 rows) |

vllm/models/deepseek_v41/nvidia/specialized_decode/engram/prefetch.py

@@ -2,15 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Engram WKV ahead of the decoder layers.
 
-With host-offloaded tables, ``Engram`` looks up its rows on a prefetch stream
-before the decoder layers. The rest of the WKV chain (head all-gather, MXFP8
-quantize, column-parallel GEMM, output all-gather) reads only those rows, so
-for decode steps ``EngramPrefetch`` runs it on a side stream right after each
-layer's lookup, and the layer only waits for it before the fused post
-kernel. The kernels, inputs and collectives are those of ``Engram.forward``.
+``Engram`` looks up its rows before the decoder layers: inline for GPU-resident
+tables, on a prefetch stream for host-offloaded ones. The rest of the WKV chain
+(head all-gather, MXFP8 quantize, column-parallel GEMM, output all-gather) reads
+only those rows, so for decode steps ``EngramPrefetch`` runs it on a side stream
+right after each layer's lookup, and the layer only waits for it before the
+fused post kernel. The kernels, inputs and collectives are those of
+``Engram.forward``.
 
 The chain issues TP collectives on the side stream, so whether a step starts it
-depends only on the step's rows, which every TP rank shares.
+depends only on the step's rows, which every TP rank shares. They stay on NCCL:
+a Lamport all-gather here, spinning beside the executor's kernels, hung
+multi-request decode steps.
 """
 
 import torch
@@ -30,14 +33,12 @@
 
 
 def prefetchable(engram: Engram) -> bool:
-    """Whether ``engram``'s WKV can run ahead: host-offloaded, TP-only tables
-    and a CuTe DSL MXFP8 WKV (FlashInfer's CUTLASS MXFP8 GEMM shares one
-    process-wide workspace with GEMMs on other streams)."""
+    """Whether ``engram``'s WKV can run ahead: TP-only tables and a CuTe DSL
+    MXFP8 WKV (FlashInfer's CUTLASS MXFP8 GEMM shares one process-wide
+    workspace with GEMMs on other streams)."""
     kernel = getattr(getattr(engram.wkv, "quant_method", None), "kernel", None)
     return (
-        engram._prefetch_stream is not None
-        and engram._prefetch_done is not None
-        and engram._batch is None
+        engram._batch is None
         and engram.embed_tokens.dp_size == 1
         and not engram.use_sequence_parallel
         and isinstance(kernel, FlashInferCutedslMxfp8LinearKernel)
@@ -75,8 +76,11 @@ def start(self, num_tokens: int) -> None:
         """Start every layer's WKV chain once its lookup is done."""
         stream = self.stream
         for engram, (out, done) in self._buffers.items():
-            assert engram._prefetch_done is not None
-            stream.wait_event(engram._prefetch_done)
+            if engram._prefetch_done is None:
+                # GPU-resident rows were looked up inline on this stream.
+                stream.wait_stream(torch.cuda.current_stream())
+            else:
+                stream.wait_event(engram._prefetch_done)
             with torch.cuda.stream(stream):
                 heads = engram.staged_rows[:num_tokens]
                 if engram.embed_tokens.tp_size > 1: