diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 6fbaa06e294a..381c19145633 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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 @@ -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. @@ -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 @@ -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( @@ -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 @@ -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 []) @@ -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 @@ -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, ) @@ -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, @@ -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,