diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index bfbd82d66..e9fbb96d2 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 0c91ff41a..6488cb121 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 82ca1903a..eae96940e 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/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 ) diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index f5ef566d2..6cac59ea0 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index db0327b6a..d761a1b16 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 38636296f..766d74809 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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", ) diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 5bf76d3d5..024bee974 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/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", ) diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index c3154a5f9..661ed41c8 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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( diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index c0a149f74..6df731814 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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( diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index ef338410c..0c13137d7 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 91a9c3d83..35bf0bc4a 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/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) diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index a9cce1013..66c65580f 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/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),