dsa: widen the fp8 k-cache quant kernel's token_id to int64 (#30859)

python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py

@@ -134,7 +134,7 @@ def _dequantize_k_cache_fast_kernel(
     DIM_NOPE: tl.constexpr,
     DIM_ROPE: tl.constexpr,
 ):
-    token_id = tl.program_id(0)
+    token_id = tl.program_id(0).to(tl.int64)
     raw_block_id = tl.program_id(1)
 
     if raw_block_id < NUM_NOPE_BLOCKS:
@@ -249,8 +249,8 @@ def _dequantize_k_cache_paged_kernel(
     DIM_NOPE: tl.constexpr,
     DIM_ROPE: tl.constexpr,
 ):
-    token_id = tl.program_id(0)
-    token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
+    token_id = tl.program_id(0).to(tl.int64)
+    token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int64)
     raw_block_id = tl.program_id(1)
 
     if raw_block_id < NUM_NOPE_BLOCKS:

python/sglang/kernels/ops/attention/dsa/quant_k_cache.py

@@ -283,7 +283,7 @@ def _quantize_k_cache_fast_kernel(
     FP8_MIN: tl.constexpr,
     FP8_MAX: tl.constexpr,
 ):
-    token_id = tl.program_id(0)
+    token_id = tl.program_id(0).to(tl.int64)
     raw_block_id = tl.program_id(1)
 
     if raw_block_id < NUM_NOPE_BLOCKS:

python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py

@@ -239,7 +239,7 @@ def _dequantize_k_cache_paged_kernel(
 ):
     # One program per token: load page_table[token_id] once and emit all
     # NUM_SCALE_TILES nope tiles + rope tail via tl.static_range.
-    token_id = tl.program_id(0)
+    token_id = tl.program_id(0).to(tl.int64)
     loc = tl.load(page_table_ptr + token_id).to(tl.int64)
     page_idx = loc // PAGE_SIZE
     in_page = loc % PAGE_SIZE

python/sglang/kernels/ops/attention/dsv4/quant_k_cache.py

@@ -26,7 +26,7 @@ def _quant_k_cache_fused_kernel(
     FP8_MAX: tl.constexpr,
     EPS: tl.constexpr,
 ):
-    token_id = tl.program_id(0)
+    token_id = tl.program_id(0).to(tl.int64)
     tile_id = tl.program_id(1)
 
     if tile_id == NUM_TILES: