fix update_step_context.

Author: Reinerzhou

lmdeploy/pytorch/backends/dlinfer/maca/op_backend.py

@@ -49,10 +49,10 @@ def update_step_context(cls, step_context):
         is_unpaged_prefill = False
         q_start_loc = torch.cat((torch.tensor([0], device=device),
                                  step_context.q_seqlens.cumsum(0))).int()
-        q_seqlens_cpu = step_context.q_seqlens.cpu()
-        kv_seqlens_cpu = step_context.kv_seqlens.cpu()
-        max_q_seq_len = torch.max(q_seqlens_cpu).item()
-        max_kv_seq_len = torch.max(kv_seqlens_cpu).item()
+        q_seqlens = step_context.q_seqlens.int()
+        kv_seqlens = step_context.kv_seqlens.int()
+        max_q_seq_len = torch.max(q_seqlens).item()
+        max_kv_seq_len = torch.max(kv_seqlens).item()
 
         if not step_context.is_decoding:
             is_unpaged_prefill = \
@@ -95,10 +95,10 @@ def update_step_context(cls, step_context):
         attn_meta_cls = cls.get_attention_metadata_cls()
         attn_metadata = attn_meta_cls(
             step_context.is_decoding,
-            step_context.block_offsets,
+            step_context.block_offsets.int(),
             q_start_loc=q_start_loc,
-            q_seqlens=q_seqlens_cpu,
-            kv_seqlens=kv_seqlens_cpu,
+            q_seqlens=q_seqlens,
+            kv_seqlens=kv_seqlens,
             kv_start_indices=kv_start_indices,
             block_size=block_size,
             attention_mask=attention_mask,