Skip to content

[Perf] dsv41: Use the fused ratio-2 compressor in the prefill graph #46230

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hnyls2002 merged 3 commits into from
Oct 24, 2026
Merged

[Perf] dsv41: Use the fused ratio-2 compressor in the prefill graph #46230

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
810a2d5
run the fused ratio-2 compressor in the prefill graph
hnyls2002 Oct 24, 2026
6be8890
keep separate kv and score buffers where the compressor has no fused …
hnyls2002 Oct 24, 2026
7cf7d8a
trim comments
hnyls2002 Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 51 additions & 17 deletions python/sglang/srt/layers/attention/deepseek_v4_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
)
from sglang.srt.layers.attention.dsv4.dsv41_sparse import (
_rope_fq4,
fused_low_ratio_compress_supported,
token_req_indices,
)
from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
Expand Down Expand Up @@ -303,10 +304,13 @@ def put(name, value):
return

if run_compressor and layer.compressor is not None:
kv, score = layer.compressor.project(x[:real])
put("kv", kv)
if score is not None:
put("score", score)
if layer.compressor.use_fused_gate:
put("kv_score", layer.compressor.project_fused(x[:real]))
else:
kv, score = layer.compressor.project(x[:real])
put("kv", kv)
if score is not None:
put("score", score)
if run_indexer and layer.indexer is not None:
# Only the captured consumers read these; a full top-k layer projects
# once inside its eager break (_index_q_and_weights), so skip them there.
Expand Down Expand Up @@ -1043,6 +1047,9 @@ class DSV4Metadata:
# Shared by all low-ratio source layers; graph replay refreshes them live.
low_ratio_req_indices: Optional[torch.Tensor] = None
low_ratio_pos_i64: Optional[torch.Tensor] = None
# The prefill graph's int32 [2, request slots]: each request's first row and row
# count, for the captured ratio-2 compressor. Slots past the batch hold no rows.
low_ratio_extend_offsets: Optional[torch.Tensor] = None
# Each low ratio's index-K slots (scoring.prefill_k_slots), gathered once by the
# extend's metadata builder; the late-layer tail's metadata reads the same dict.
low_ratio_k_slots: Optional[Dict[int, torch.Tensor]] = None
Expand Down Expand Up @@ -1096,6 +1103,10 @@ def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata
maybe_copy_inplace(
self.low_ratio_pos_i64, src=static_metadata.low_ratio_pos_i64
)
maybe_copy_inplace(
self.low_ratio_extend_offsets,
src=static_metadata.low_ratio_extend_offsets,
)
self.low_ratio_k_slots = static_metadata.low_ratio_k_slots
self.low_ratio_index_plans = static_metadata.low_ratio_index_plans
maybe_copy_inplace(
Expand Down Expand Up @@ -1677,6 +1688,17 @@ def create(compress_ratio: Literal[4, 128]):
metadata.low_ratio_pos_i64 = core_attn_metadata.positions_casual.to(
torch.int64
)
# One slot per request the bucket can hold. DP padding extends
# extend_seq_lens past the requests; those rows belong to no slot.
num_reqs = extend_start_loc.shape[0]
offsets = torch.zeros(
(2, min(self.req_to_token.shape[0], padded_num_tokens)),
dtype=torch.int32,
device=extend_seq_lens.device,
)
offsets[0, :num_reqs] = extend_start_loc
offsets[1, :num_reqs] = extend_seq_lens[:num_reqs]
metadata.low_ratio_extend_offsets = offsets
return metadata

@property
Expand Down Expand Up @@ -3207,13 +3229,14 @@ def _source_projection_buffers(self, num_tokens: int, ratio: int) -> dict:
bufs = {
"q": zeros(num_tokens, heads, dim, dtype=torch.bfloat16),
"w": zeros(num_tokens, heads, dtype=torch.bfloat16),
"kv": zeros(
num_tokens,
latent,
dtype=torch.bfloat16 if ratio == 1 else torch.float32,
),
}
if ratio == 2:
if ratio == 1:
bufs["kv"] = zeros(num_tokens, latent, dtype=torch.bfloat16)
elif fused_low_ratio_compress_supported():
# The fused compressor reads | kv | score | rows.
bufs["kv_score"] = zeros(num_tokens, 2 * latent, dtype=torch.float32)
else:
bufs["kv"] = zeros(num_tokens, latent, dtype=torch.float32)
bufs["score"] = zeros(num_tokens, latent, dtype=torch.float32)
# Keep the previous allocations alive: captured graphs still hold them.
sets[ratio] = [bufs] + (have or [])
Expand Down Expand Up @@ -3411,9 +3434,19 @@ def forward_low_ratio_sources(
layer, x, q_lora, pos, bufs, run_compressor, captured_index
)
if run_compressor and layer.compressor is not None:
self._low_ratio_compress_torch(
layer, x, req, pos, projected=(bufs["kv"], bufs.get("score"))
)
if layer.compressor.use_fused_gate:
self._low_ratio_compress_fused(
layer,
x,
req,
pos,
projected=bufs["kv_score"],
extend_offsets=meta.low_ratio_extend_offsets.unbind(),
)
else:
self._low_ratio_compress_torch(
layer, x, req, pos, projected=(bufs["kv"], bufs.get("score"))
)
if run_indexer and layer.indexer is not None:
if captured_index:
# Captured: sparse scores of the published blocks only, so
Expand Down Expand Up @@ -3580,10 +3613,11 @@ def _low_ratio_compress_decode(self, layer, x, req, pos) -> None:
)

def _low_ratio_compress_fused(
self, layer, x, req, pos, *, draft_len=1, extend_offsets=None
self, layer, x, req, pos, *, draft_len=1, extend_offsets=None, projected=None
) -> None:
"""Decode or target-verify; an extend batch passes ``extend_offsets``, the int32
(first row, row count) of each request."""
(first row, row count) of each request. ``projected`` is the compressor's
projection of ``x`` when the caller already has it."""
from sglang.kernels.ops.attention.dsv4.fp4_indexer_rope import (
index_k_norm_rope_pack_store,
)
Expand Down Expand Up @@ -3612,7 +3646,7 @@ def _low_ratio_compress_fused(
if layer.compress_ratio == 1:
# At ratio 1, c1_out_loc equals the int64 raw_out_loc supplied by the scheduler.
latent = c1_decode_norm_rope_store(
compressor.wkv(x),
compressor.wkv(x) if projected is None else projected,
compressor.norm.weight.data,
pos,
core.raw_out_loc,
Expand All @@ -3628,7 +3662,7 @@ def _low_ratio_compress_fused(
# KVAndScore rows use | kv | score |, addressed as req * ring_size + pos % ring_size.
state = pool.get_attention_compress_states(layer_id)
args = (
compressor.project_fused(x),
compressor.project_fused(x) if projected is None else projected,
state.kv_score_buffer.kv_score,
compressor.norm.weight.data,
pos,
Expand Down
Loading