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

[DSV4] Use int64 for compressor out_loc tensors (#27973)

Browse Source
This commit is contained in:
Cheng Wan
2026-06-11 17:45:34 -07:00
committed by GitHub
parent 97a0031799
commit 2e74ff192c
4 changed files with 19 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+14 -14
View File
@@ -44,7 +44,7 @@ struct FusedNormRopeStoreParams {
const void* __restrict__ handle; // plan decode / compress const void* __restrict__ handle; // plan decode / compress
const void* __restrict__ weight; const void* __restrict__ weight;
const float* __restrict__ freqs_cis; const float* __restrict__ freqs_cis;
const int32_t* __restrict__ out_loc; const int64_t* __restrict__ out_loc;
uint8_t* __restrict__ kvcache; uint8_t* __restrict__ kvcache;
float eps; float eps;
uint32_t compress_ratio; uint32_t compress_ratio;
@@ -90,7 +90,7 @@ INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRop
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim; const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
int32_t position; int32_t position;
int32_t out_loc; int64_t out_loc;
if constexpr (kMode == CompressExtend) { if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id]; const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
if (plan.is_invalid()) return; if (plan.is_invalid()) return;
@@ -204,8 +204,8 @@ INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRop
const auto abs_max = warp::reduce_max(local_max); const auto abs_max = warp::reduce_max(local_max);
const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX; const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
const auto inv_scale = 1.0f / scale; const auto inv_scale = 1.0f / scale;
const int32_t page = out_loc >> kPageBits; const int64_t page = out_loc >> kPageBits;
const int32_t offset = out_loc & ((1 << kPageBits) - 1); const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes; const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 128; const auto value_ptr = page_ptr + offset * 128;
const auto scale_ptr = page_ptr + (128 << kPageBits) + offset * 4; const auto scale_ptr = page_ptr + (128 << kPageBits) + offset * 4;
@@ -244,7 +244,7 @@ INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNor
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim; const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
int32_t position; int32_t position;
int32_t out_loc; int64_t out_loc;
if constexpr (kMode == CompressExtend) { if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id]; const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
if (plan.is_invalid()) return; if (plan.is_invalid()) return;
@@ -351,8 +351,8 @@ INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNor
const uint8_t packed1 = quant_fp4_e2m1(data[2] * inv_scale) | (quant_fp4_e2m1(data[3] * inv_scale) << 4); const uint8_t packed1 = quant_fp4_e2m1(data[2] * inv_scale) | (quant_fp4_e2m1(data[3] * inv_scale) << 4);
const uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8); const uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8);
const int32_t page = out_loc >> kPageBits; const int64_t page = out_loc >> kPageBits;
const int32_t offset = out_loc & ((1 << kPageBits) - 1); const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes; const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 64; const auto value_ptr = page_ptr + offset * 64;
const auto scale_ptr = page_ptr + (64 << kPageBits) + offset * 4; const auto scale_ptr = page_ptr + (64 << kPageBits) + offset * 4;
@@ -398,7 +398,7 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim; const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
int32_t position; int32_t position;
int32_t out_loc; int64_t out_loc;
if constexpr (kMode == CompressExtend) { if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id]; const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
if (plan.is_invalid()) return; if (plan.is_invalid()) return;
@@ -450,8 +450,8 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
} }
} }
const int32_t page = out_loc >> kPageBits; const int64_t page = out_loc >> kPageBits;
const int32_t offset = out_loc & ((1 << kPageBits) - 1); const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes; const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576); const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
@@ -560,7 +560,7 @@ struct FusedNormRopeKernel {
.with_device(device_) .with_device(device_)
.verify(freqs_cis); .verify(freqs_cis);
TensorMatcher({-1}) // out_loc TensorMatcher({-1}) // out_loc
.with_dtype<int32_t>() .with_dtype<int64_t>()
.with_device(device_) .with_device(device_)
.verify(out_loc); .verify(out_loc);
TensorMatcher({-1, -1}) // cache TensorMatcher({-1, -1}) // cache
@@ -587,7 +587,7 @@ struct FusedNormRopeKernel {
.handle = plan.data_ptr(), .handle = plan.data_ptr(),
.weight = weight.data_ptr(), .weight = weight.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()), .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.out_loc = static_cast<const int32_t*>(out_loc.data_ptr()), .out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()), .kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.eps = eps, .eps = eps,
.compress_ratio = compress_ratio, .compress_ratio = compress_ratio,
@@ -625,7 +625,7 @@ struct FusedNormRopeKernel {
TensorMatcher({N, kHeadDim}).with_dtype<DType>().with_device(device_).verify(input); TensorMatcher({N, kHeadDim}).with_dtype<DType>().with_device(device_).verify(input);
TensorMatcher({kHeadDim}).with_dtype<DType>().with_device(device_).verify(weight); TensorMatcher({kHeadDim}).with_dtype<DType>().with_device(device_).verify(weight);
TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis); TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis);
TensorMatcher({-1}).with_dtype<int32_t>().with_device(device_).verify(out_loc); TensorMatcher({-1}).with_dtype<int64_t>().with_device(device_).verify(out_loc);
TensorMatcher({-1, -1}).with_strides({kFp4PageBytes, 1}).with_dtype<uint8_t>().with_device(device_).verify(kvcache); TensorMatcher({-1, -1}).with_strides({kFp4PageBytes, 1}).with_dtype<uint8_t>().with_device(device_).verify(kvcache);
switch (mode) { switch (mode) {
@@ -646,7 +646,7 @@ struct FusedNormRopeKernel {
.handle = plan.data_ptr(), .handle = plan.data_ptr(),
.weight = weight.data_ptr(), .weight = weight.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()), .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.out_loc = static_cast<const int32_t*>(out_loc.data_ptr()), .out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()), .kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.eps = eps, .eps = eps,
.compress_ratio = compress_ratio, .compress_ratio = compress_ratio,
python/sglang/srt/layers/attention/dsv4/compressor_v2.py
+2 -5
View File
@@ -519,12 +519,9 @@ class CompressorBackendMixin:
kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id) kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
page_size = token_to_kv_pool.get_extra_key_page_size(layer_id) page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"): if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
# The v2 compressor writes directly into the raw C4 KV tensor. out_loc = compress_kv_pool._translate_loc_to_hisparse_device(
# HiSparse C4 therefore needs the physical C4 location here.
# The compress kernel requires an int32 write location.
out_loc = compress_kv_pool.translate_loc_to_hisparse_device(
out_loc out_loc
).to(torch.int32) )
self._forward_compress_all_in_one( self._forward_compress_all_in_one(
kv_score_buffer=state_pool.kv_score_buffer.kv_score, kv_score_buffer=state_pool.kv_score_buffer.kv_score,
kv_score_input=kv_score_input, kv_score_input=kv_score_input,
python/sglang/srt/layers/attention/dsv4/metadata_kernel.py
+2 -2
View File
@@ -107,12 +107,12 @@ def _init_compressed_attn_metadata_triton(
bs = seq_lens.shape[0] bs = seq_lens.shape[0]
device = seq_lens.device device = seq_lens.device
c4_out_loc = torch.empty(bs, dtype=torch.int32, device=device) c4_out_loc = torch.empty(bs, dtype=torch.int64, device=device)
c4_positions = torch.empty(bs, dtype=torch.int32, device=device) c4_positions = torch.empty(bs, dtype=torch.int32, device=device)
c4_seq_lens_raw = torch.empty(bs, dtype=torch.int32, device=device) c4_seq_lens_raw = torch.empty(bs, dtype=torch.int32, device=device)
c4_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device) c4_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
c128_out_loc = torch.empty(bs, dtype=torch.int32, device=device) c128_out_loc = torch.empty(bs, dtype=torch.int64, device=device)
c128_positions = torch.empty(bs, dtype=torch.int32, device=device) c128_positions = torch.empty(bs, dtype=torch.int32, device=device)
c128_seq_lens_raw = torch.empty(bs, dtype=torch.int32, device=device) c128_seq_lens_raw = torch.empty(bs, dtype=torch.int32, device=device)
c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device) c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
test/registered/jit/deepseek_v4/test_fp4_indexer.py
+1 -1
View File
@@ -148,7 +148,7 @@ def test_fp4_fused_norm_rope_store_layout(num_tokens: int) -> None:
plan = CompressorDecodePlan.generate_legacy( plan = CompressorDecodePlan.generate_legacy(
compress_ratio, req_pool_indices, seq_lens compress_ratio, req_pool_indices, seq_lens
) )
loc = torch.arange(num_tokens, device="cuda", dtype=torch.int32) loc = torch.arange(num_tokens, device="cuda", dtype=torch.int64)
freqs_cis = precompute_freqs_cis( freqs_cis = precompute_freqs_cis(
64, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1 64, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1
).to("cuda") ).to("cuda")