[Fix] Use int64 seq_lens across all CUDA graph runners and backends (#27840)

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -1535,6 +1535,8 @@ class DeepseekV4AttnBackend(
     ) -> DSV4AttnMetadata:
         assert self.swa_page_size == SWA_WINDOW
 
+        seq_lens_casual = seq_lens_casual.to(torch.int32)
+
         swa_page_indices = self.get_swa_page_indices(
             seq_lens_casual=seq_lens_casual,
             req_pool_indices_repeated=req_pool_indices_repeated,

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -1426,6 +1426,8 @@ class DeepseekV4HipRadixBackend(
     ) -> DSV4AttnMetadata:
         assert self.swa_page_size == SWA_WINDOW
 
+        seq_lens_casual = seq_lens_casual.to(torch.int32)
+
         swa_page_indices = self.get_swa_page_indices(
             seq_lens_casual=seq_lens_casual,
             req_pool_indices_repeated=req_pool_indices_repeated,

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -472,7 +472,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         """
         if forward_mode.is_decode_or_idle():
             assert seq_lens_cpu is not None
-            kv_len_arr_cpu = seq_lens_cpu[:bs]
+            kv_len_arr_cpu = seq_lens_cpu[:bs].to(torch.int32)
             self.cuda_graph_kv_indptr_cpu[1 : bs + 1] = torch.cumsum(
                 kv_len_arr_cpu, dim=0
             )

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -570,7 +570,7 @@ def build_decode_registry(
         GraphSlot(
             "seq_lens",
             _bs,
-            torch.int32,
+            torch.int64,
             axis="bs",
             padding_policy=PaddingPolicy.FILL_SENTINEL,
             pad_value=seq_len_fill_value,
@@ -578,7 +578,7 @@ def build_decode_registry(
         GraphSlot(
             "seq_lens_cpu",
             _bs,
-            torch.int32,
+            torch.int64,
             axis="bs",
             device=torch.device("cpu"),
             padding_policy=PaddingPolicy.FILL_SENTINEL,

python/sglang/srt/model_executor/model_runner.py

@@ -3212,7 +3212,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         forward_batch: ForwardBatch,
         pp_proxy_tensors=None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
-        if not self.server_args.enable_pdmux and self.device == "cuda":
+        if not self.server_args.enable_pdmux:
             forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors)
         # Set extra arguments
         pdmux_override = False
@@ -3302,7 +3302,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 ret = self.prefill_cuda_graph_runner.replay(forward_batch, **kwargs)
             return (ret, can_run_graph)
 
-        if not self.server_args.enable_pdmux and self.device == "cuda":
+        if not self.server_args.enable_pdmux:
             forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors)
 
         # Launch model forward
@@ -3363,7 +3363,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # called from the idle path can re-read a prior batch's req_pool
         # indices and trigger SWA mapping use-after-free.
         if forward_batch.batch_size > 0:
-            if not self.server_args.enable_pdmux and self.device == "cuda":
+            if not self.server_args.enable_pdmux:
                 forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors)
             self.attn_backend.init_forward_metadata(forward_batch)
         else:

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -197,7 +197,7 @@ def _allocate_decode_buffers(
         input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
         input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
         req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
-        seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
+        seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int64)
         out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
         positions = torch.zeros((max_num_token,), dtype=torch.int64)
         mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
@@ -269,7 +269,7 @@ def _allocate_decode_buffers(
     seq_lens_cpu = torch.full(
         (max_bs,),
         seq_len_fill_value,
-        dtype=torch.int32,
+        dtype=torch.int64,
         device="cpu",
     )
 

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -99,7 +99,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
             input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
             input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
             req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
-            seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
+            seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int64)
             out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
             out_cache_loc_swa = (
                 torch.zeros((max_num_token,), dtype=torch.int64)
@@ -177,7 +177,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
         seq_lens_cpu = torch.full(
             (max_bs,),
             seq_len_fill_value,
-            dtype=torch.int32,
+            dtype=torch.int64,
             device="cpu",
         )
 

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -170,7 +170,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
                 else None
             )
             seq_lens = torch.full(
-                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
+                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
             )
             extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
             topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
@@ -204,7 +204,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
                 global_num_tokens_for_logprob_gpu = None
 
         seq_lens_cpu = torch.full(
-            (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32, device="cpu"
+            (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
         )
 
         self.buffers = EagleDraftInputBuffers(

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -167,7 +167,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
             )
             seq_lens = torch.full(
-                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
+                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
             )
             extend_seq_lens = torch.full(
                 (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32
@@ -221,7 +221,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             )
 
         seq_lens_cpu = torch.full(
-            (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32, device="cpu"
+            (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
         )
 
         self.buffers = EagleDraftExtendInputBuffers(

python/sglang/srt/speculative/eagle_info.py

@@ -117,7 +117,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
             spec_steps=spec_steps,
             capture_hidden_mode=CaptureHiddenMode.FULL,
             seq_lens_sum=0,
-            seq_lens_cpu=torch.empty((0,), dtype=torch.int32),
+            seq_lens_cpu=torch.empty((0,), dtype=torch.int64),
         )
 
     def prepare_for_verify(self, batch: ScheduleBatch, page_size: int):
@@ -935,8 +935,8 @@ class EagleDraftExtendInput(SpecInput):
             num_accept_tokens=torch.empty((0,), device=device, dtype=torch.int32),
             num_accept_tokens_cpu=[],
             input_ids=torch.empty((0,), device=device, dtype=torch.long),
-            seq_lens=torch.empty((0,), device=device, dtype=torch.int32),
-            seq_lens_cpu=torch.empty((0,), dtype=torch.int32),
+            seq_lens=torch.empty((0,), device=device, dtype=torch.int64),
+            seq_lens_cpu=torch.empty((0,), dtype=torch.int64),
             req_pool_indices=torch.empty((0,), device=device, dtype=torch.int64),
             capture_hidden_mode=capture_hidden_mode,
         )

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -98,7 +98,7 @@ class FrozenKVMTPCudaGraphRunner:
             self.draft_attn_backend.get_cuda_graph_seq_len_fill_value()
         )
         seq_lens_cpu = torch.full(
-            (self.max_num_token,), self.seq_len_fill_value, dtype=torch.int32
+            (self.max_num_token,), self.seq_len_fill_value, dtype=torch.int64
         )
 
         if self.enable_torch_compile:
@@ -109,7 +109,7 @@ class FrozenKVMTPCudaGraphRunner:
             positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
             seq_lens = torch.full(
-                (self.max_num_token,), self.seq_len_fill_value, dtype=torch.int32
+                (self.max_num_token,), self.seq_len_fill_value, dtype=torch.int64
             )
             topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
             topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64)

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -819,8 +819,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.tensor([0, 1], dtype=torch.int64),
             out_cache_loc=torch.tensor([100, 101], dtype=torch.int64),
             req_pool_indices=torch.tensor([1, 2], dtype=torch.int64),
-            seq_lens=torch.tensor([7, 8], dtype=torch.int32),
-            seq_lens_cpu=torch.tensor([7, 8], dtype=torch.int32),
+            seq_lens=torch.tensor([7, 8], dtype=torch.int64),
+            seq_lens_cpu=torch.tensor([7, 8], dtype=torch.int64),
             mrope_positions=torch.tensor([[0, 1], [0, 1], [0, 1]], dtype=torch.int64),
         )
         # Poison tails so resets are observable.
@@ -847,14 +847,16 @@ class TestBuildDecodeRegistry(unittest.TestCase):
         self.assertTrue(torch.equal(rp[2:4], torch.tensor([0, 0])))
         # FILL_SENTINEL: head copied, tail = seq_len_fill_value.
         sl = reg.get_slot("seq_lens").buffer
-        self.assertTrue(torch.equal(sl[:2], torch.tensor([7, 8], dtype=torch.int32)))
+        self.assertEqual(sl.dtype, torch.int64)
+        self.assertTrue(torch.equal(sl[:2], torch.tensor([7, 8], dtype=torch.int64)))
         self.assertTrue(
-            torch.equal(sl[2:4], torch.tensor([FILL, FILL], dtype=torch.int32))
+            torch.equal(sl[2:4], torch.tensor([FILL, FILL], dtype=torch.int64))
         )
         slc = reg.get_slot("seq_lens_cpu").buffer
         self.assertEqual(slc.device.type, "cpu")
+        self.assertEqual(slc.dtype, torch.int64)
         self.assertTrue(
-            torch.equal(slc[2:4], torch.tensor([FILL, FILL], dtype=torch.int32))
+            torch.equal(slc[2:4], torch.tensor([FILL, FILL], dtype=torch.int64))
         )
         # 2D mrope via slice_fn.
         mr = reg.get_slot("mrope_positions").buffer
@@ -879,8 +881,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.zeros(8, dtype=torch.int64),
             out_cache_loc=torch.zeros(8, dtype=torch.int64),
             req_pool_indices=torch.zeros(4, dtype=torch.int64),
-            seq_lens=torch.full((4,), 5, dtype=torch.int32),
-            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int32),
+            seq_lens=torch.full((4,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int64),
             mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
             global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
             global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),
@@ -915,8 +917,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.zeros(8, dtype=torch.int64),
             out_cache_loc=torch.zeros(8, dtype=torch.int64),
             req_pool_indices=torch.zeros(4, dtype=torch.int64),
-            seq_lens=torch.full((4,), 5, dtype=torch.int32),
-            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int32),
+            seq_lens=torch.full((4,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int64),
             mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
             num_token_non_padded=ntnp,
             global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
@@ -976,8 +978,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.arange(2, dtype=torch.int64),
             out_cache_loc=torch.arange(2, dtype=torch.int64),
             req_pool_indices=torch.zeros(2, dtype=torch.int64),
-            seq_lens=torch.full((2,), 5, dtype=torch.int32),
-            seq_lens_cpu=torch.full((2,), 5, dtype=torch.int32),
+            seq_lens=torch.full((2,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((2,), 5, dtype=torch.int64),
             global_num_tokens_gpu=gnt,
             global_num_tokens_for_logprob_gpu=gntlp,
         )
@@ -1014,8 +1016,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.zeros(8, dtype=torch.int64),
             out_cache_loc=torch.zeros(8, dtype=torch.int64),
             req_pool_indices=torch.zeros(4, dtype=torch.int64),
-            seq_lens=torch.full((4,), 5, dtype=torch.int32),
-            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int32),
+            seq_lens=torch.full((4,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int64),
             mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
             global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
             global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),
@@ -1062,8 +1064,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.zeros(8, dtype=torch.int64),
             out_cache_loc=torch.zeros(8, dtype=torch.int64),
             req_pool_indices=torch.zeros(4, dtype=torch.int64),
-            seq_lens=torch.full((4,), 5, dtype=torch.int32),
-            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int32),
+            seq_lens=torch.full((4,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int64),
             mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
             global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
             global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),
@@ -1110,8 +1112,8 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             positions=torch.zeros(8, dtype=torch.int64),
             out_cache_loc=torch.zeros(8, dtype=torch.int64),
             req_pool_indices=torch.zeros(4, dtype=torch.int64),
-            seq_lens=torch.full((4,), 5, dtype=torch.int32),
-            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int32),
+            seq_lens=torch.full((4,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int64),
             mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
             global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
             global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),