cuda graph: adjust capture time num-non-padded-tokens to align capture with replay (#22404)

2026-04-11 10:27:50 +08:00
parent 0b4f5c9fcb
commit 58f863956c
1 changed files with 13 additions and 0 deletions
@@ -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: