From 2e74ff192c86f246be1a28c5a4cdb61dec907be7 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:45:34 -0700 Subject: [PATCH] [DSV4] Use int64 for compressor out_loc tensors (#27973) --- .../csrc/deepseek_v4/fused_norm_rope_v2.cuh | 28 +++++++++---------- .../layers/attention/dsv4/compressor_v2.py | 7 ++--- .../layers/attention/dsv4/metadata_kernel.py | 4 +-- .../jit/deepseek_v4/test_fp4_indexer.py | 2 +- 4 files changed, 19 insertions(+), 22 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh index 32dd8c8ee..0e789263c 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh @@ -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(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(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(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(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(packed0) | (static_cast(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(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(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() + .with_dtype() .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(freqs_cis.data_ptr()), - .out_loc = static_cast(out_loc.data_ptr()), + .out_loc = static_cast(out_loc.data_ptr()), .kvcache = static_cast(kvcache.data_ptr()), .eps = eps, .compress_ratio = compress_ratio, @@ -625,7 +625,7 @@ struct FusedNormRopeKernel { TensorMatcher({N, kHeadDim}).with_dtype().with_device(device_).verify(input); TensorMatcher({kHeadDim}).with_dtype().with_device(device_).verify(weight); TensorMatcher({-1, kRopeDim}).with_dtype().with_device(device_).verify(freqs_cis); - TensorMatcher({-1}).with_dtype().with_device(device_).verify(out_loc); + TensorMatcher({-1}).with_dtype().with_device(device_).verify(out_loc); TensorMatcher({-1, -1}).with_strides({kFp4PageBytes, 1}).with_dtype().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(freqs_cis.data_ptr()), - .out_loc = static_cast(out_loc.data_ptr()), + .out_loc = static_cast(out_loc.data_ptr()), .kvcache = static_cast(kvcache.data_ptr()), .eps = eps, .compress_ratio = compress_ratio, diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index a499e429d..c3caaf86f 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -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, diff --git a/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py b/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py index 34cbcde08..5e4b999f7 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py @@ -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) diff --git a/test/registered/jit/deepseek_v4/test_fp4_indexer.py b/test/registered/jit/deepseek_v4/test_fp4_indexer.py index e56e8326c..36dbfd811 100644 --- a/test/registered/jit/deepseek_v4/test_fp4_indexer.py +++ b/test/registered/jit/deepseek_v4/test_fp4_indexer.py @@ -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")