From 58f863956c4879a60be7fb0a494784311b438efa Mon Sep 17 00:00:00 2001
From: Yilong Zhao <74357408+happierpig@users.noreply.github.com>
Date: Fri, 10 Apr 2026 19:27:50 -0700
Subject: [PATCH] cuda graph: adjust capture time num-non-padded-tokens to
 align capture with replay (#22404)

---
 .../sglang/srt/model_executor/cuda_graph_runner.py  | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py
index e906de2b7..b4e21bb3b 100644
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -870,7 +870,20 @@ class CudaGraphRunner:
             encoder_lens = None
         mrope_positions = buffers.mrope_positions[:, :num_tokens]
         next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens]
+
+        # Adjust for attention TP if needed (matching replay path in
+        # populate_from_forward_batch).
         buffers.num_token_non_padded[...] = num_tokens
+        if (
+            enable_num_token_non_padded(self.model_runner.server_args)
+            and self.require_gathered_buffer
+            and not self.nsa_enable_prefill_cp
+        ):
+            local = compute_local_num_token_non_padded(
+                global_num_token_non_padded=buffers.num_token_non_padded,
+                num_tokens_per_dp=num_tokens,
+            )
+            buffers.num_token_non_padded.copy_(local)
 
         # pipeline parallelism
         if self.pp_size > 1: