minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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__ weight;
const float* __restrict__ freqs_cis;
const int32_t* __restrict__ out_loc;
const int64_t* __restrict__ out_loc;
uint8_t* __restrict__ kvcache;
float eps;
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;
int32_t position;
int32_t out_loc;
int64_t out_loc;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
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 scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
const auto inv_scale = 1.0f / scale;
const int32_t page = out_loc >> kPageBits;
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 128;
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;
int32_t position;
int32_t out_loc;
int64_t out_loc;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
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 uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8);
const int32_t page = out_loc >> kPageBits;
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 64;
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;
int32_t position;
int32_t out_loc;
int64_t out_loc;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
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 int32_t offset = out_loc & ((1 << kPageBits) - 1);
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
@@ -560,7 +560,7 @@ struct FusedNormRopeKernel {
.with_device(device_)
.verify(freqs_cis);
TensorMatcher({-1}) // out_loc
.with_dtype<int32_t>()
.with_dtype<int64_t>()
.with_device(device_)
.verify(out_loc);
TensorMatcher({-1, -1}) // cache
@@ -587,7 +587,7 @@ struct FusedNormRopeKernel {
.handle = plan.data_ptr(),
.weight = weight.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()),
.eps = eps,
.compress_ratio = compress_ratio,
@@ -625,7 +625,7 @@ struct FusedNormRopeKernel {
TensorMatcher({N, kHeadDim}).with_dtype<DType>().with_device(device_).verify(input);
TensorMatcher({kHeadDim}).with_dtype<DType>().with_device(device_).verify(weight);
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);
switch (mode) {
@@ -646,7 +646,7 @@ struct FusedNormRopeKernel {
.handle = plan.data_ptr(),
.weight = weight.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()),
.eps = eps,
.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)
page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
# The v2 compressor writes directly into the raw C4 KV tensor.
# 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 = compress_kv_pool._translate_loc_to_hisparse_device(
out_loc
).to(torch.int32)
)
self._forward_compress_all_in_one(
kv_score_buffer=state_pool.kv_score_buffer.kv_score,
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]
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_seq_lens_raw = 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_seq_lens_raw = 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(
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(
64, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1
).to("cuda")