perf(dsv4): tile C128 compression to reduce register spills

lightllm/models/deepseek_v4/layer_infer/compressor.py

@@ -126,48 +126,62 @@ def fused_compress(
     block_state = triton.next_power_of_2(candidate_dim)
     block_head = triton.next_power_of_2(head_dim)
 
-    _fused_compress_norm_rope_kernel[(kv_score.shape[0],)](
-        kv_score,
-        kv_score.stride(0),
-        kv_score.stride(1),
-        req_to_compressor_state,
-        infer_state.position_ids,
-        infer_state._dsv4_token_to_batch_idx,
-        infer_state.b_req_idx,
-        infer_state.b_mtp_index,
-        infer_state.b_seq_len,
-        infer_state.b_ready_cache_len,
-        infer_state.b_q_start_loc,
-        compressed_kv_slots,
-        norm_weight,
-        eps,
-        cos_table,
-        cos_table.stride(0),
-        cos_table.stride(1),
-        sin_table,
-        sin_table.stride(0),
-        sin_table.stride(1),
-        out_buffer,
-        HEAD_DIM=head_dim,
-        CANDIDATE_DIM=candidate_dim,
-        STATE_DIM=state_dim,
-        COMPRESS_RATIO=compress_ratio,
-        WINDOW_SIZE=state.config.window_size,
-        IS_C4=is_c4,
-        IS_PREFILL=infer_state.is_prefill,
-        STATE_CAPACITY_TOKENS=state_capacity_tokens,
-        ROPE_HEAD_DIM=qk_rope_head_dim,
-        FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
-        AMAX_MIN=1e-4,
-        NOPE_DIM=head_dim - qk_rope_head_dim,
-        QUANT_BLOCK=64,
-        SCALE_BYTES=(head_dim - qk_rope_head_dim) // 64 + 1,
-        PAGE_SIZE=out_page_size,
-        BYTES_PER_PAGE=out_buffer.stride(0),
-        BLOCK=block_head,
-        IS_IN_INDEXER=is_in_indexer,
-        num_warps=4,
-    )
+    scratch = None
+    stages = ("compress_norm_rope_pack",)
+    if compress_ratio == 128:
+        scratch = infer_state.dsv4_workspace.get_c128_compress_scratch(kv_score.shape[0], infer_state.microbatch_index)
+        stages = ("window_compress", "norm_rope_pack")
+    for stage in stages:
+        store_compressed = stage == "window_compress"
+        load_compressed = stage == "norm_rope_pack"
+        block = 64 if store_compressed else block_head
+        warps = 8 if store_compressed else 4
+        grid = (kv_score.shape[0], triton.cdiv(head_dim, block)) if store_compressed else (kv_score.shape[0],)
+        _fused_compress_norm_rope_kernel[grid](
+            kv_score,
+            kv_score.stride(0),
+            kv_score.stride(1),
+            req_to_compressor_state,
+            infer_state.position_ids,
+            infer_state._dsv4_token_to_batch_idx,
+            infer_state.b_req_idx,
+            infer_state.b_mtp_index,
+            infer_state.b_seq_len,
+            infer_state.b_ready_cache_len,
+            infer_state.b_q_start_loc,
+            compressed_kv_slots,
+            norm_weight,
+            eps,
+            cos_table,
+            cos_table.stride(0),
+            cos_table.stride(1),
+            sin_table,
+            sin_table.stride(0),
+            sin_table.stride(1),
+            out_buffer,
+            HEAD_DIM=head_dim,
+            CANDIDATE_DIM=candidate_dim,
+            STATE_DIM=state_dim,
+            COMPRESS_RATIO=compress_ratio,
+            WINDOW_SIZE=state.config.window_size,
+            IS_C4=is_c4,
+            IS_PREFILL=infer_state.is_prefill,
+            STATE_CAPACITY_TOKENS=state_capacity_tokens,
+            ROPE_HEAD_DIM=qk_rope_head_dim,
+            FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
+            AMAX_MIN=1e-4,
+            NOPE_DIM=head_dim - qk_rope_head_dim,
+            QUANT_BLOCK=64,
+            SCALE_BYTES=(head_dim - qk_rope_head_dim) // 64 + 1,
+            PAGE_SIZE=out_page_size,
+            BYTES_PER_PAGE=out_buffer.stride(0),
+            BLOCK=block,
+            IS_IN_INDEXER=is_in_indexer,
+            compress_scratch=scratch,
+            STORE_COMPRESSED=store_compressed,
+            LOAD_COMPRESSED=load_compressed,
+            num_warps=warps,
+        )
 
     _save_compressor_state_kernel[(kv_score.shape[0],)](
         kv_score,

lightllm/models/deepseek_v4/model.py

@@ -240,6 +240,7 @@ def _init_custom(self):
         if os.getenv("LIGHTLLM_DSV4_PREFILL_OVERLAP", "1") == "1" and not self.args.enable_prefill_microbatch_overlap:
             prefill_aux_stream = torch.cuda.Stream()
         self.dsv4_workspace = DeepseekV4Workspace(self)
+        self.dsv4_workspace.init_c128_compress_scratch(self)
         if prefill_aux_stream is not None and self.run_mode == "prefill" and 4 in self.config["compress_ratios"]:
             with torch.cuda.stream(prefill_aux_stream):
                 self.dsv4_workspace.init_c4_prefill_aux(self)

lightllm/models/deepseek_v4/triton_kernel/compressor.py

@@ -110,72 +110,85 @@ def _fused_compress_norm_rope_kernel(
     BYTES_PER_PAGE: tl.constexpr,
     BLOCK: tl.constexpr,
     IS_IN_INDEXER: tl.constexpr,
+    compress_scratch,
+    STORE_COMPRESSED: tl.constexpr = False,
+    LOAD_COMPRESSED: tl.constexpr = False,
 ):
     token_idx = tl.program_id(0)
     out_slot = tl.load(compressed_kv_slots + token_idx).to(tl.int64)
     if out_slot < 0:
         return
 
     position = tl.load(position_ids + token_idx)
-    batch_idx = tl.load(token_to_batch_idx + token_idx) if IS_PREFILL else token_idx
-    req_idx = tl.load(b_req_idx + batch_idx).to(tl.int64)
-    seq_len = tl.load(b_seq_len + batch_idx)
-    if IS_PREFILL:
-        ready_len = tl.load(b_ready_cache_len + batch_idx)
-        q_start = tl.load(b_q_start_loc + batch_idx)
+    if LOAD_COMPRESSED:
+        offs = tl.arange(0, BLOCK)
+        dim_mask = offs < HEAD_DIM
+        compressed_kv = tl.load(compress_scratch + token_idx * HEAD_DIM + offs, mask=dim_mask, other=0.0)
     else:
-        mtp_index = tl.load(b_mtp_index + batch_idx)
-        ready_len = position - mtp_index
-        q_start = token_idx - mtp_index
+        batch_idx = tl.load(token_to_batch_idx + token_idx) if IS_PREFILL else token_idx
+        req_idx = tl.load(b_req_idx + batch_idx).to(tl.int64)
+        seq_len = tl.load(b_seq_len + batch_idx)
+        if IS_PREFILL:
+            ready_len = tl.load(b_ready_cache_len + batch_idx)
+            q_start = tl.load(b_q_start_loc + batch_idx)
+        else:
+            mtp_index = tl.load(b_mtp_index + batch_idx)
+            ready_len = position - mtp_index
+            q_start = token_idx - mtp_index
 
-    token_offsets = tl.arange(0, WINDOW_SIZE)
-    start = position - WINDOW_SIZE + 1
-    gather_pos = start + token_offsets
-    valid_pos = (gather_pos >= 0) & (gather_pos < seq_len)
-    use_current = (gather_pos >= ready_len) & valid_pos
-    current_idx = q_start + (gather_pos - ready_len)
-    state_valid = valid_pos & ~use_current
-    state_slot = req_idx * STATE_CAPACITY_TOKENS + gather_pos % STATE_CAPACITY_TOKENS
-    if IS_C4:
-        head_offset = tl.where(token_offsets >= COMPRESS_RATIO, HEAD_DIM, 0)
-    else:
-        head_offset = token_offsets * 0
+        token_offsets = tl.arange(0, WINDOW_SIZE)
+        start = position - WINDOW_SIZE + 1
+        gather_pos = start + token_offsets
+        valid_pos = (gather_pos >= 0) & (gather_pos < seq_len)
+        use_current = (gather_pos >= ready_len) & valid_pos
+        current_idx = q_start + (gather_pos - ready_len)
+        state_valid = valid_pos & ~use_current
+        state_slot = req_idx * STATE_CAPACITY_TOKENS + gather_pos % STATE_CAPACITY_TOKENS
+        if IS_C4:
+            head_offset = tl.where(token_offsets >= COMPRESS_RATIO, HEAD_DIM, 0)
+        else:
+            head_offset = token_offsets * 0
 
-    offs = tl.arange(0, BLOCK)
-    dim_mask = offs < HEAD_DIM
-    current_mask = use_current[:, None] & dim_mask[None, :]
-    state_mask = state_valid[:, None] & dim_mask[None, :]
+        offs = tl.arange(0, BLOCK) + tl.program_id(1) * BLOCK
+        dim_mask = offs < HEAD_DIM
+        current_mask = use_current[:, None] & dim_mask[None, :]
+        state_mask = state_valid[:, None] & dim_mask[None, :]
 
-    cur_ptrs = (
-        kv_score + current_idx[:, None] * kv_score_stride0 + (head_offset[:, None] + offs[None, :]) * kv_score_stride1
-    )
-    cur_kv = tl.load(
-        cur_ptrs,
-        mask=current_mask,
-        other=0.0,
-    )
-    cur_score = tl.load(
-        cur_ptrs + CANDIDATE_DIM * kv_score_stride1,
-        mask=current_mask,
-        other=float("-inf"),
-    )
+        cur_ptrs = (
+            kv_score
+            + current_idx[:, None] * kv_score_stride0
+            + (head_offset[:, None] + offs[None, :]) * kv_score_stride1
+        )
+        cur_kv = tl.load(
+            cur_ptrs,
+            mask=current_mask,
+            other=0.0,
+        )
+        cur_score = tl.load(
+            cur_ptrs + CANDIDATE_DIM * kv_score_stride1,
+            mask=current_mask,
+            other=float("-inf"),
+        )
 
-    state_ptrs = req_to_compressor_state + state_slot[:, None] * STATE_DIM + head_offset[:, None] + offs[None, :]
-    state_kv = tl.load(
-        state_ptrs,
-        mask=state_mask,
-        other=0.0,
-    )
-    state_score = tl.load(
-        state_ptrs + CANDIDATE_DIM,
-        mask=state_mask,
-        other=float("-inf"),
-    )
+        state_ptrs = req_to_compressor_state + state_slot[:, None] * STATE_DIM + head_offset[:, None] + offs[None, :]
+        state_kv = tl.load(
+            state_ptrs,
+            mask=state_mask,
+            other=0.0,
+        )
+        state_score = tl.load(
+            state_ptrs + CANDIDATE_DIM,
+            mask=state_mask,
+            other=float("-inf"),
+        )
 
-    kv = tl.where(current_mask, cur_kv, state_kv)
-    score = tl.where(current_mask, cur_score, state_score)
-    score = tl.softmax(score, dim=0)
-    compressed_kv = tl.sum(kv * score, axis=0)
+        kv = tl.where(current_mask, cur_kv, state_kv)
+        score = tl.where(current_mask, cur_score, state_score)
+        score = tl.softmax(score, dim=0)
+        compressed_kv = tl.sum(kv * score, axis=0)
+        if STORE_COMPRESSED:
+            tl.store(compress_scratch + token_idx * HEAD_DIM + offs, compressed_kv, mask=dim_mask)
+            return
     # rms_norm
     rms_w = tl.load(norm_weight + offs, mask=dim_mask, other=0.0)
     variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_DIM

lightllm/models/deepseek_v4/workspace.py

@@ -65,6 +65,30 @@ def __init__(self, model):
         self.token_capacity = ((int(model.batch_max_tokens) + tp - 1) // tp) * tp
         self.flashmla_prefill_q = None
         self.c4_prefill_aux = None
+        self.c128_compress_scratch = None
+
+    def init_c128_compress_scratch(self, model):
+        regions = 1 + int(model.args.enable_decode_microbatch_overlap or model.args.enable_prefill_microbatch_overlap)
+        decode_requests = model.max_req_num
+        if model.args.enable_decode_microbatch_overlap:
+            decode_requests = (decode_requests + 1) // 2
+        decode_tokens = decode_requests * model.mtp_manager.get_decode_tokens_per_request(model.is_mtp_draft_model)
+        # Decode budgets count requests, while the compressor sees expanded MTP
+        # rows after TP/SP and graph padding. Keep prefill-only arenas unchanged.
+        token_capacity = max(self.token_capacity, model._align_decode_batch_size(max(1, decode_tokens)))
+        if not model.disable_cudagraph:
+            token_capacity = max(token_capacity, model.graph_max_batch_size)
+        self.c128_compress_scratch = torch.empty(
+            (regions, token_capacity, model.config["head_dim"]), dtype=torch.float32, device="cuda"
+        )
+
+    def get_c128_compress_scratch(self, token_num: int, microbatch_index: int):
+        regions, capacity, _ = self.c128_compress_scratch.shape
+        if not 0 <= microbatch_index < regions:
+            raise ValueError(f"DeepSeek-V4 C128 workspace has {regions} regions, got region {microbatch_index}")
+        if not 0 <= token_num <= capacity:
+            raise ValueError(f"DeepSeek-V4 C128 workspace capacity {capacity} cannot cover {token_num} tokens")
+        return self.c128_compress_scratch[microbatch_index, :token_num]
 
     def init_c4_prefill_aux(self, model):
         self.c4_prefill_aux = DeepseekV4C4PrefillWorkspace(