Use int64 as indices for set_kv_buffer (#3039)

2025-01-21 19:46:09 -08:00
parent a42213dbd4
commit 3d8f1c9bcf
6 changed files with 30 additions and 37 deletions
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/python/sglang/srt/layers/logits_processor.py
@@ -296,7 +296,7 @@ def fused_softcap_kernel(
    n_elements,
    BLOCK_SIZE: tl.constexpr,
 ):
-    pid = tl.program_id(0)
+    pid = tl.program_id(0).to(tl.int64)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements