minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xinyuan Tong
2026-08-26 14:22:24 -07:00
committed by GitHub
parent 2935bb8e79
commit a8d716ce8d
4 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py
+3 -3
View File
@@ -134,7 +134,7 @@ def _dequantize_k_cache_fast_kernel(
DIM_NOPE: tl.constexpr, DIM_NOPE: tl.constexpr,
DIM_ROPE: 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) raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS: if raw_block_id < NUM_NOPE_BLOCKS:
@@ -249,8 +249,8 @@ def _dequantize_k_cache_paged_kernel(
DIM_NOPE: tl.constexpr, DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr, DIM_ROPE: tl.constexpr,
): ):
token_id = tl.program_id(0) token_id = tl.program_id(0).to(tl.int64)
token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32) token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int64)
raw_block_id = tl.program_id(1) raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS: if raw_block_id < NUM_NOPE_BLOCKS:
python/sglang/kernels/ops/attention/dsa/quant_k_cache.py
+1 -1
View File
@@ -283,7 +283,7 @@ def _quantize_k_cache_fast_kernel(
FP8_MIN: tl.constexpr, FP8_MIN: tl.constexpr,
FP8_MAX: 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) raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS: if raw_block_id < NUM_NOPE_BLOCKS:
python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py
+1 -1
View File
@@ -239,7 +239,7 @@ def _dequantize_k_cache_paged_kernel(
): ):
# One program per token: load page_table[token_id] once and emit all # One program per token: load page_table[token_id] once and emit all
# NUM_SCALE_TILES nope tiles + rope tail via tl.static_range. # 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) loc = tl.load(page_table_ptr + token_id).to(tl.int64)
page_idx = loc // PAGE_SIZE page_idx = loc // PAGE_SIZE
in_page = loc % PAGE_SIZE in_page = loc % PAGE_SIZE
python/sglang/kernels/ops/attention/dsv4/quant_k_cache.py
+1 -1
View File
@@ -26,7 +26,7 @@ def _quant_k_cache_fused_kernel(
FP8_MAX: tl.constexpr, FP8_MAX: tl.constexpr,
EPS: 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) tile_id = tl.program_id(1)
if tile_id == NUM_TILES: if tile_id == NUM_TILES: