[Perf] [NVIDIA] dspark: Move the host work between DSv4.1 decode graphs into the graphs for the megakernel path

python/sglang/kernels/ops/embeddings/engram_hash.py

@@ -28,12 +28,19 @@ def _engram_commit_history_kernel(
     tokens_ptr,
     slots_ptr,
     commit_ptr,
+    live_ptr,
+    gate_ptr,
     HISTORY_STRIDE: tl.constexpr,
     HISTORY_WIDTH: tl.constexpr,
     TOKEN_STRIDE: tl.constexpr,
     BLOCK: tl.constexpr,
+    GATED: tl.constexpr,
 ):
     row = tl.program_id(0)
+    if GATED:
+        # Captured rows past the live batch, or a step the graph does not commit.
+        if (tl.load(live_ptr + row) == 0) | (tl.load(gate_ptr) == 0):
+            return
     slot = tl.load(slots_ptr + row).to(tl.int64)
     commit = tl.load(commit_ptr + row)
     col = tl.arange(0, BLOCK)
@@ -60,8 +67,12 @@ def engram_commit_history(
     verify_ids: torch.Tensor,
     req_slots: torch.Tensor,
     commit_lens: torch.Tensor,
+    live_rows: Optional[torch.Tensor] = None,
+    gate: Optional[torch.Tensor] = None,
 ) -> None:
-    """Update distinct live request slots with anchor + correct drafts, not bonus."""
+    """Update distinct live request slots with anchor + correct drafts, not bonus.
+    With ``live_rows`` and ``gate``, rows whose live entry is 0 are skipped, and all
+    rows when ``gate[0]`` is 0."""
     bs = req_slots.numel()
     width = history.shape[1]
     if bs == 0 or width == 0:
@@ -73,10 +84,13 @@ def engram_commit_history(
         verify_ids,
         req_slots,
         commit_lens,
+        live_rows if live_rows is not None else commit_lens,
+        gate if gate is not None else commit_lens,
         HISTORY_STRIDE=history.stride(0),
         HISTORY_WIDTH=width,
         TOKEN_STRIDE=verify_ids.stride(0),
         BLOCK=triton.next_power_of_2(width),
+        GATED=live_rows is not None,
         num_warps=4,
     )
 

python/sglang/kernels/ops/speculative/cache_locs.py

@@ -546,6 +546,10 @@ def assign_extend_cache_locs_uniform(
     PAGE_SIZE: tl.constexpr = 1,
     num_v_pages=0,
     num_swa_pages=0,
+    positions_out=None,
+    anchor_src=None,
+    anchor_out=None,
+    anchor_stride=0,
 ):
     """Uniform-length variant of assign_extend_cache_locs: every row extends
     exactly draft_token_num tokens, so the end offset is start +
@@ -556,6 +560,8 @@ def assign_extend_cache_locs_uniform(
     kv_start = tl.load(start_offset + pid)
     token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
     out_cache_ptr = out_cache_loc + pid * draft_token_num
+    if anchor_out is not None:
+        tl.store(anchor_out + pid * anchor_stride, tl.load(anchor_src + pid))
 
     offs = tl.arange(0, BLOCK_SIZE)
     num_loop = tl.cdiv(draft_token_num, BLOCK_SIZE)
@@ -564,6 +570,12 @@ def assign_extend_cache_locs_uniform(
         mask = o < draft_token_num
         data = tl.load(token_pool + kv_start + o, mask=mask)
         tl.store(out_cache_ptr + o, data, mask=mask)
+        if positions_out is not None:
+            tl.store(
+                positions_out + pid * draft_token_num + o,
+                kv_start.to(tl.int64) + o,
+                mask=mask,
+            )
         if full_v2p is not None:
             page = data // PAGE_SIZE
             valid = mask & (data >= 0) & (page < num_v_pages)
@@ -594,6 +606,9 @@ def assign_extend_cache_locs_uniform_func(
     full_out: Optional[torch.Tensor] = None,
     swa_out: Optional[torch.Tensor] = None,
     page_size: int = 1,
+    positions_out: Optional[torch.Tensor] = None,
+    anchor_src: Optional[torch.Tensor] = None,
+    anchor_out: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """assign_extend_cache_locs for the uniform case (all rows extend exactly
     draft_token_num tokens, e.g. spec target-verify prep). Computes end
@@ -602,6 +617,8 @@ def assign_extend_cache_locs_uniform_func(
     and SWA write windows in the same launch; the returned ids stay virtual.
     Mapping tables and output buffers must be contiguous, with one output
     per gathered token. Translation is rank-local (no DCP widening).
+    ``positions_out`` ([batch, draft_token_num] int64) gets start + column and
+    ``anchor_out[:, 0]`` gets ``anchor_src``, in the same launch.
     """
     if full_v2p is not None:
         assert full_out is not None and full_out.numel() == batch_size * draft_token_num
@@ -632,11 +649,22 @@ def assign_extend_cache_locs_uniform_func(
             PAGE_SIZE=page_size,
             num_v_pages=0 if full_v2p is None else full_v2p.numel(),
             num_swa_pages=0 if swa_v2p is None else swa_v2p.numel(),
+            positions_out=positions_out,
+            anchor_src=anchor_src,
+            anchor_out=anchor_out,
+            anchor_stride=0 if anchor_out is None else anchor_out.stride(0),
         )
         return out_cache_loc
 
     # NPU / CPU platforms: fall back to the end_offset-tensor path.
     assert full_v2p is None, "fused write translation requires a Triton device"
+    if positions_out is not None:
+        positions_out.copy_(
+            start_offset.to(torch.int64).view(-1, 1)
+            + torch.arange(draft_token_num, device=start_offset.device)
+        )
+    if anchor_out is not None:
+        anchor_out[:, 0].copy_(anchor_src.view(-1))
     return assign_extend_cache_locs_func(
         req_pool_indices=req_pool_indices,
         req_to_token=req_to_token,