From 90e3fbf0c308dc0f24e44873f954c441039d0d76 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Thu, 7 May 2026 09:37:19 +0800 Subject: [PATCH 01/10] support mla chunk-prefill & prefix_cache --- ...get_position_ids_and_mask_encoder_batch.cu | 37 +- .../layers/attention/mla_attention_backend.py | 469 +++++++++++++++++- .../model_executor/models/deepseek_v3.py | 76 ++- 3 files changed, 552 insertions(+), 30 deletions(-) diff --git a/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu b/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu index 63bc77c9afc..1aeb85c4496 100644 --- a/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu +++ b/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu @@ -17,7 +17,8 @@ __global__ void GetPositionIdsAndMaskEncoderBatchKernel( const int* seq_lens_encoder, // [bsz] 每个批次的 encoder 长度 - const int* seq_lens_decoder, // [bsz] 每个批次的 decoder 长度 + const int* seq_lens_decoder, // [bsz] 每个批次的 decoder 长度 (对于 MLA + // prefix cache，这是 cached KV 长度) const int* seq_lens_this_time, int* position_ids, // 输出的一维 position_ids const int bsz) { // 批次大小 @@ -25,11 +26,16 @@ __global__ void GetPositionIdsAndMaskEncoderBatchKernel( int tid = threadIdx.x; if (tid >= bsz) return; - // 动态计算当前批次的偏移量 + // 动态计算当前批次的偏移量。 + // 每个 batch 只会贡献 encoder_len 或 seq_lens_this_time 中的一个， + // 而非两者之和（chunked prefill 时 encoder_len > 0 与 decoder_len > 0 + // 同时成立， + // 但该 batch 只有 encoder_len 个真实 token）。 int offset = 0; for (int i = 0; i < tid; i++) { - offset += seq_lens_encoder[i]; - if (seq_lens_decoder[i] > 0) { + if (seq_lens_encoder[i] > 0) { + offset += seq_lens_encoder[i]; + } else if (seq_lens_decoder[i] > 0) { offset += seq_lens_this_time[i]; } } @@ -39,16 +45,21 @@ __global__ void GetPositionIdsAndMaskEncoderBatchKernel( int decoder_len = seq_lens_decoder[tid]; int seq_len_this_time = seq_lens_this_time[tid]; - // 写入 encoder 的 position_ids - for (int i = 0; i < encoder_len; i++) { - position_ids[offset + i] = i; - } - offset += encoder_len; - - // 写入 decoder 的 position_ids - if (decoder_len > 0) { + // For MLA with prefix cache support: + // - Chunked prefill (encoder_len > 0 && decoder_len > 0): + // only writes encoder positions starting at decoder_len (cached length). + // - First-chunk prefill (encoder_len > 0 && decoder_len == 0): + // writes encoder positions starting at 0. + // - Pure decode (encoder_len == 0 && decoder_len > 0): + // writes seq_lens_this_time decode positions starting at decoder_len. + if (encoder_len > 0) { + int start_pos = (decoder_len > 0) ? decoder_len : 0; + for (int i = 0; i < encoder_len; i++) { + position_ids[offset + i] = start_pos + i; + } + } else if (decoder_len > 0) { for (int i = 0; i < seq_len_this_time; i++) { - position_ids[offset + i] = decoder_len + i; // 使用 decoder 长度本身 + position_ids[offset + i] = decoder_len + i; } } } diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 7e2e1066f0b..153b2290250 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -67,6 +67,284 @@ ) from fastdeploy.spec_decode import SpecMethod +# ============================================================================ +# Latent Cache Read Kernel for Prefix Cache Support +# ============================================================================ + + +@enable_compat_on_triton_kernel +@triton.jit() +def read_latent_from_cache_kernel( + latent_cache, + block_tables, + cache_kv_lens, + cu_seqlens_cached_kv, + output_kv_c, + output_k_pe, + block_size: tl.constexpr, + kv_lora_rank: tl.constexpr, + qk_rope_head_dim: tl.constexpr, + LATENT_DIM: tl.constexpr, +): + """ + Kernel to read latent vectors (kv_c and k_pe) from paged latent cache. + Each program instance handles one cached token. + + Args: + latent_cache: [num_blocks, 1, block_size, kv_lora_rank + qk_rope_head_dim] + block_tables: [batch_size, max_blocks_per_seq] + cache_kv_lens: [batch_size] - cached KV length for each request + cu_seqlens_cached_kv: [batch_size + 1] - cumulative sequence lengths for cached KV + output_kv_c: [total_cached_tokens, kv_lora_rank] + output_k_pe: [total_cached_tokens, qk_rope_head_dim] + """ + # Global token index in the output + token_idx = tl.program_id(axis=0) + + # Find which batch this token belongs to using binary search on cu_seqlens + # For simplicity, we use a linear scan (could be optimized with binary search) + batch_id = 0 + for i in range(cu_seqlens_cached_kv.shape[0] - 1): + if token_idx >= tl.load(cu_seqlens_cached_kv + i) and token_idx < tl.load(cu_seqlens_cached_kv + i + 1): + batch_id = i + break + + # Local token index within the batch + cu_start = tl.load(cu_seqlens_cached_kv + batch_id) + local_token_idx = token_idx - cu_start + + # Get the physical block and offset + block_idx = local_token_idx // block_size + block_offset = local_token_idx % block_size + + # Get physical block id from block_tables + physical_block_id = tl.load(block_tables + batch_id * block_tables.shape[1] + block_idx) + + # Load latent vector from cache + # latent_cache shape: [num_blocks, 1, block_size, kv_lora_rank + qk_rope_head_dim] + latent_base = latent_cache + physical_block_id * LATENT_DIM * block_size + block_offset * LATENT_DIM + + # Read kv_c (first kv_lora_rank dimensions) + kv_c_offsets = tl.arange(0, kv_lora_rank) + kv_c_value = tl.load(latent_base + kv_c_offsets, mask=kv_c_offsets < kv_lora_rank) + + # Read k_pe (last qk_rope_head_dim dimensions) + k_pe_offsets = tl.arange(kv_lora_rank, kv_lora_rank + qk_rope_head_dim) + k_pe_value = tl.load(latent_base + k_pe_offsets, mask=k_pe_offsets < LATENT_DIM) + + # Store outputs + output_kv_c_base = output_kv_c + token_idx * kv_lora_rank + tl.store(output_kv_c_base + kv_c_offsets, kv_c_value, mask=kv_c_offsets < kv_lora_rank) + + output_k_pe_base = output_k_pe + token_idx * qk_rope_head_dim + k_pe_out_offsets = tl.arange(0, qk_rope_head_dim) + tl.store(output_k_pe_base + k_pe_out_offsets, k_pe_value) + + +def read_latent_from_cache( + latent_cache: paddle.Tensor, + block_tables: paddle.Tensor, + cache_kv_lens: paddle.Tensor, + cu_seqlens_cached_kv: paddle.Tensor, + total_cached_tokens: int, + block_size: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + """ + Read latent vectors (kv_c and k_pe) from paged latent cache for prefix cache support. + + Args: + latent_cache: [num_blocks, 1, block_size, kv_lora_rank + qk_rope_head_dim] + block_tables: [batch_size, max_blocks_per_seq] + cache_kv_lens: [batch_size] - cached KV length for each request + cu_seqlens_cached_kv: [batch_size + 1] - cumulative sequence lengths for cached KV + total_cached_tokens: Total number of cached tokens across all requests + block_size: Block size for paged attention + kv_lora_rank: LoRA rank for KV compression + qk_rope_head_dim: Dimension of RoPE part in key + + Returns: + cached_kv_c: [total_cached_tokens, kv_lora_rank] + cached_k_pe: [total_cached_tokens, qk_rope_head_dim] + """ + if total_cached_tokens == 0: + return None, None + + # latent_dim = kv_lora_rank + qk_rope_head_dim + + # Allocate output tensors + cached_kv_c = paddle.empty([total_cached_tokens, kv_lora_rank], dtype=latent_cache.dtype) + cached_k_pe = paddle.empty([total_cached_tokens, qk_rope_head_dim], dtype=latent_cache.dtype) + + # Use a simpler approach: iterate through each batch and gather latent vectors + # IMPORTANT: Only process batches that have prefix cache (use cu_seqlens_cached_kv) + bsz = cu_seqlens_cached_kv.shape[0] - 1 + output_idx = 0 + + print(f"[DEBUG read_latent_from_cache] total_cached_tokens={total_cached_tokens}, bsz={bsz}") + print(f"[DEBUG read_latent_from_cache] cu_seqlens_cached_kv={cu_seqlens_cached_kv.tolist()}") + print(f"[DEBUG read_latent_from_cache] block_tables shape={block_tables.shape}") + + for batch_id in range(bsz): + # Get the number of cached tokens for this batch from cu_seqlens_cached_kv + cu_start = ( + cu_seqlens_cached_kv[batch_id].item() + if hasattr(cu_seqlens_cached_kv[batch_id], "item") + else cu_seqlens_cached_kv[batch_id] + ) + cu_end = ( + cu_seqlens_cached_kv[batch_id + 1].item() + if hasattr(cu_seqlens_cached_kv[batch_id + 1], "item") + else cu_seqlens_cached_kv[batch_id + 1] + ) + cache_len = cu_end - cu_start + + if cache_len <= 0: + continue + + # Debug: Print cache reading info + print(f"[DEBUG read_latent_from_cache] batch_id={batch_id}, cache_len={cache_len}") + + # Read tokens from multiple blocks if cache_len > block_size + local_idx = 0 + while local_idx < cache_len: + block_idx = local_idx // block_size + block_offset = local_idx % block_size + tokens_to_read = min(block_size - block_offset, cache_len - local_idx) + + physical_block_id = block_tables[batch_id, block_idx].item() + + # Debug: Print block access info + print( + f"[DEBUG read_latent_from_cache] block_idx={block_idx}, block_offset={block_offset}, physical_block_id={physical_block_id}" + ) + + # Load latent vectors from this block + for offset in range(tokens_to_read): + latent_vec = latent_cache[physical_block_id, 0, block_offset + offset, :] + + # Split into kv_c and k_pe + cached_kv_c[output_idx] = latent_vec[:kv_lora_rank] + cached_k_pe[output_idx] = latent_vec[kv_lora_rank:] + output_idx += 1 + + local_idx += tokens_to_read + + print(f"[DEBUG read_latent_from_cache] Total cached tokens read: {output_idx}") + return cached_kv_c, cached_k_pe + + +def interleave_cached_and_new_latent( + cached_kv_c: paddle.Tensor, + cached_k_pe: paddle.Tensor, + new_compressed_kv: paddle.Tensor, + new_k_pe: paddle.Tensor, + cu_seqlens_cached_kv: paddle.Tensor, + seq_lens_encoder: paddle.Tensor, + cu_seqlens_q: paddle.Tensor, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + """ + Interleave cached and new latent vectors per batch for correct FlashAttention layout. + + The key insight is that FlashAttention expects tokens to be ordered by batch: + [batch0_tokens, batch1_tokens, ...] + where each batch's tokens should be [cached_tokens, new_tokens]. + + Args: + cached_kv_c: [total_cached_tokens, kv_lora_rank] - cached KV latent (from read_latent_from_cache) + cached_k_pe: [total_cached_tokens, qk_rope_head_dim] - cached K RoPE (from read_latent_from_cache) + new_compressed_kv: [total_new_tokens, kv_lora_rank] - new KV latent from current forward + new_k_pe: [total_new_tokens, qk_rope_head_dim] - new K RoPE from current forward (already with RoPE applied) + cu_seqlens_cached_kv: [batch_size + 1] - cumulative sequence lengths for cached KV + seq_lens_encoder: [batch_size] - number of new tokens per batch + cu_seqlens_q: [batch_size + 1] - cumulative sequence lengths for new tokens + kv_lora_rank: LoRA rank for KV compression + qk_rope_head_dim: Dimension of RoPE part in key + + Returns: + full_compressed_kv: [total_tokens, kv_lora_rank] - properly interleaved KV latent + full_k_pe: [total_tokens, qk_rope_head_dim] - properly interleaved K RoPE + """ + bsz = cu_seqlens_cached_kv.shape[0] - 1 + + # Calculate total output size + total_cached = cached_kv_c.shape[0] if cached_kv_c is not None and cached_kv_c.numel() > 0 else 0 + total_new = new_compressed_kv.shape[0] + total_tokens = total_cached + total_new + + # Allocate output tensors + full_compressed_kv = paddle.empty([total_tokens, kv_lora_rank], dtype=new_compressed_kv.dtype) + full_k_pe = paddle.empty([total_tokens, qk_rope_head_dim], dtype=new_k_pe.dtype) + + # Track indices into cached and new tensors + cached_idx = 0 + new_idx = 0 + out_position = 0 # Track output position for each batch + + print( + f"[DEBUG interleave_cached_and_new_latent] bsz={bsz}, total_cached={total_cached}, total_new={total_new}, total_tokens={total_tokens}" + ) + + for batch_id in range(bsz): + # Number of cached tokens for this batch + cu_cached_start = ( + cu_seqlens_cached_kv[batch_id].item() + if hasattr(cu_seqlens_cached_kv[batch_id], "item") + else cu_seqlens_cached_kv[batch_id] + ) + cu_cached_end = ( + cu_seqlens_cached_kv[batch_id + 1].item() + if hasattr(cu_seqlens_cached_kv[batch_id + 1], "item") + else cu_seqlens_cached_kv[batch_id + 1] + ) + num_cached = cu_cached_end - cu_cached_start + + # Number of new tokens for this batch + cu_new_start = ( + cu_seqlens_q[batch_id].item() if hasattr(cu_seqlens_q[batch_id], "item") else cu_seqlens_q[batch_id] + ) + cu_new_end = ( + cu_seqlens_q[batch_id + 1].item() + if hasattr(cu_seqlens_q[batch_id + 1], "item") + else cu_seqlens_q[batch_id + 1] + ) + num_new = cu_new_end - cu_new_start + + print( + f"[DEBUG interleave] batch_id={batch_id}, num_cached={num_cached}, num_new={num_new}, cached_idx={cached_idx}, out_position={out_position}" + ) + + # Output position for this batch (sequential, no gaps) + out_start = out_position + + # Copy cached tokens first (if any) + if num_cached > 0 and cached_kv_c is not None: + full_compressed_kv[out_start : out_start + num_cached] = cached_kv_c[cached_idx : cached_idx + num_cached] + full_k_pe[out_start : out_start + num_cached] = cached_k_pe[cached_idx : cached_idx + num_cached] + cached_idx += num_cached + + # Then copy new tokens + if num_new > 0: + full_compressed_kv[out_start + num_cached : out_start + num_cached + num_new] = new_compressed_kv[ + new_idx : new_idx + num_new + ] + full_k_pe[out_start + num_cached : out_start + num_cached + num_new] = new_k_pe[ + new_idx : new_idx + num_new + ] + new_idx += num_new + + # Update output position for next batch + out_position += num_cached + num_new + + print(f"[DEBUG interleave] Final: cached_idx={cached_idx}, new_idx={new_idx}, out_position={out_position}") + return full_compressed_kv, full_k_pe + + +# ============================================================================ + @enable_compat_on_triton_kernel @triton.jit() @@ -232,6 +510,20 @@ class MLAAttentionMetadata(AttentionMetadata): max_dec_len_this_time: Optional[paddle.Tensor] = None max_kv_len_this_time: Optional[paddle.Tensor] = None + # For prefix cache and chunked prefill support + # Indicates whether any request has prefix cache (cached KV from previous requests) + has_prefix_cache: bool = False + # Total number of cached KV tokens across all requests with prefix cache + total_cached_kv_tokens: int = 0 + # cu_seqlens for cached KV tokens (similar to cu_seqlens_k but for cached portion) + cu_seqlens_cached_kv: Optional[paddle.Tensor] = None + # Maximum cached KV length across all requests + max_cached_kv_len: int = 0 + # cu_seqlens_k that includes cached tokens (for FlashAttention when prefix cache is present) + cu_seqlens_k_with_cache: Optional[paddle.Tensor] = None + # Maximum total KV length (cached + new) across all requests for FlashAttention + max_total_kv_len: int = 0 + class MLAAttentionBackend(AttentionBackend): """ @@ -365,6 +657,94 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): metadata.max_dec_len_this_time = forward_meta.max_len_tensor_cpu[2] metadata.max_kv_len_this_time = forward_meta.max_len_tensor_cpu[5] + # Compute prefix cache metadata + # For MLA, prefix cache is indicated by seq_lens_decoder > 0 when there are encoder tokens + # This means the request has cached KV from a previous prefix + bsz = forward_meta.seq_lens_this_time.shape[0] + has_prefix_cache = False + total_cached_kv_tokens = 0 + max_cached_kv_len = 0 + max_total_kv_len = 0 # max(dec_len + enc_len) for FlashAttention + + # Check if any request has prefix cache + # Prefix cache exists when seq_lens_decoder > 0 + # seq_lens_decoder stores the cached KV length for chunked prefill/prefix cache + for i in range(bsz): + enc_len = ( + forward_meta.seq_lens_encoder[i].item() + if hasattr(forward_meta.seq_lens_encoder[i], "item") + else forward_meta.seq_lens_encoder[i] + ) + dec_len = ( + forward_meta.seq_lens_decoder[i].item() + if hasattr(forward_meta.seq_lens_decoder[i], "item") + else forward_meta.seq_lens_decoder[i] + ) + # Count cached tokens whenever dec_len > 0, which includes: + # - chunked-prefill continuation (enc_len>0, dec_len>0) + # - pure decode (enc_len==0, dec_len>0) — its KV is read via this path too so that + # cu_seqlens_k_with_cache stays consistent with cu_seqlens_q for the prefill + # FlashAttention call; decode batch outputs here are discarded by merge_prefill_decode_output. + if dec_len > 0: + has_prefix_cache = True + total_cached_kv_tokens += dec_len + max_cached_kv_len = max(max_cached_kv_len, dec_len) + max_total_kv_len = max(max_total_kv_len, dec_len + enc_len) + if enc_len > 0: + max_total_kv_len = max(max_total_kv_len, max(enc_len, dec_len + enc_len)) + + metadata.has_prefix_cache = has_prefix_cache + metadata.total_cached_kv_tokens = total_cached_kv_tokens + metadata.max_cached_kv_len = max_cached_kv_len + metadata.max_total_kv_len = max_total_kv_len + + # Compute cu_seqlens_cached_kv if there's prefix cache + if has_prefix_cache and total_cached_kv_tokens > 0: + cu_seqlens_cached_kv = paddle.zeros([bsz + 1], dtype=paddle.int32) + cu_seqlens_k_with_cache = paddle.zeros([bsz + 1], dtype=paddle.int32) + cumsum_cached = 0 + cumsum_total = 0 + # cu_seqlens layout must stay CONSISTENT with the input tensor layout used + # by the prefill FlashAttention call. The input `compressed_kv`/`key_pe` contains + # ALL tokens in the batch (prefill + decode), in seq_lens_this_time order. + # + # For each batch i, the interleaved layout writes: [cached_tokens, new_tokens_of_this_batch] + # - new_tokens_of_this_batch = seq_lens_this_time[i] (== enc_len for prefill, == 1 for decode) + # - cached_tokens = dec_len if dec_len > 0 else 0 + # cu_seqlens_k_with_cache must reflect this sum per batch. + # cu_seqlens_cached_kv tracks only the cached portion for read_latent_from_cache(). + for i in range(bsz): + enc_len = ( + forward_meta.seq_lens_encoder[i].item() + if hasattr(forward_meta.seq_lens_encoder[i], "item") + else forward_meta.seq_lens_encoder[i] + ) + dec_len = ( + forward_meta.seq_lens_decoder[i].item() + if hasattr(forward_meta.seq_lens_decoder[i], "item") + else forward_meta.seq_lens_decoder[i] + ) + seq_this_time = ( + forward_meta.seq_lens_this_time[i].item() + if hasattr(forward_meta.seq_lens_this_time[i], "item") + else forward_meta.seq_lens_this_time[i] + ) + print( + f"[DEBUG init_attn_meta] batch {i}: enc_len={enc_len}, dec_len={dec_len}, seq_this={seq_this_time}, cumsum_cached={cumsum_cached}, cumsum_total={cumsum_total}" + ) + if dec_len > 0: + cumsum_cached += dec_len + cumsum_total += dec_len + # Add this batch's new tokens to cumsum_total. Use seq_lens_this_time to cover + # both prefill (== enc_len) and decode (== 1) correctly. + cumsum_total += seq_this_time + cu_seqlens_cached_kv[i + 1] = cumsum_cached + cu_seqlens_k_with_cache[i + 1] = cumsum_total + print(f"[DEBUG init_attn_meta] Final cu_seqlens_cached_kv: {cu_seqlens_cached_kv.tolist()}") + print(f"[DEBUG init_attn_meta] Final cu_seqlens_k_with_cache: {cu_seqlens_k_with_cache.tolist()}") + metadata.cu_seqlens_cached_kv = cu_seqlens_cached_kv + metadata.cu_seqlens_k_with_cache = cu_seqlens_k_with_cache + # pd_disaggregation metadata.kv_signal_data_list = [None] * self.num_layers if self.pd_disaggregation_mode == "per_chunk": @@ -411,7 +791,13 @@ def forward_extend( forward_meta: ForwardMeta, ) -> paddle.Tensor: """ - Prefill阶段的前向传播 + Prefill阶段的前向传播，支持 prefix cache + + 对于 MLA 模型的 prefix cache 支持： + 1. 如果存在 prefix cache (metadata.has_prefix_cache = True) + - k 和 v 应该已经包含了 cached KV 和 new KV 的拼接 + - cu_seqlens_k 应该已经调整为包含 cached tokens + 2. 如果不存在 prefix cache，行为与之前相同 """ metadata = self.attention_metadata @@ -423,7 +809,7 @@ def forward_extend( latent_cache = forward_meta.caches[layer.layer_id] if hasattr(forward_meta, "caches") else None - # 写入缓存 + # 写入新的 KV 到缓存 (只写入新 tokens，不写入 cached 部分) prefill_mla_write_cache( compressed_kv, k_pe, @@ -439,14 +825,28 @@ def forward_extend( ) # Flash注意力计算 + # 对于 prefix cache 场景： + # - k 和 v 应该已经包含了 cached + new tokens + # - cu_seqlens_k 应该已经调整 + # - max_seqlen_k 应该是 cached_len + new_len 的最大值 + + # 获取正确的 cu_seqlens_k 和 max_seqlen_k + if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: + # When prefix cache is present, use cu_seqlens_k that includes cached tokens + cu_seqlens_k = metadata.cu_seqlens_k_with_cache + max_seqlen_k = metadata.max_total_kv_len # max(dec_len + enc_len) + else: + cu_seqlens_k = forward_meta.cu_seqlens_k + max_seqlen_k = metadata.max_enc_len_this_time + fmha_out = self.flash_attn_func( q, k, v, forward_meta.cu_seqlens_q, - forward_meta.cu_seqlens_k, - metadata.max_enc_len_this_time, - metadata.max_enc_len_this_time, + cu_seqlens_k, + metadata.max_enc_len_this_time, # max_seqlen_q - only new tokens + max_seqlen_k, # max_seqlen_k - may include cached tokens causal=self.causal, **self.flash_attn_kwargs, )[0] @@ -553,7 +953,11 @@ def forward_mixed( forward_meta: ForwardMeta, ) -> paddle.Tensor: """ - Mixed模式的前向传播 + Mixed模式的前向传播，支持 prefix cache + + 对于 MLA 模型的 prefix cache 支持： + 1. Prefill 分支：k 和 v 应该已包含 cached + new tokens + 2. Decode 分支：保持原有 latent attention 逻辑 """ metadata = self.attention_metadata speculate_decoder = self.speculative_method is not None @@ -567,7 +971,41 @@ def forward_mixed( latent_cache = forward_meta.caches[layer.layer_id] if hasattr(forward_meta, "caches") else None + # Prefill branch: k is not None if k is not None: + # Debug: Verify tensor shapes and sequence lengths + bsz = forward_meta.cu_seqlens_q.shape[0] - 1 + total_q_tokens = q.shape[0] + total_k_tokens = k.shape[0] + + # Calculate expected cu_seqlens_k_with_cache + if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: + expected_k_len = ( + metadata.cu_seqlens_k_with_cache[bsz].item() + if hasattr(metadata.cu_seqlens_k_with_cache[bsz], "item") + else metadata.cu_seqlens_k_with_cache[bsz] + ) + else: + expected_k_len = ( + forward_meta.cu_seqlens_k[bsz].item() + if hasattr(forward_meta.cu_seqlens_k[bsz], "item") + else forward_meta.cu_seqlens_k[bsz] + ) + + # Debug output + print( + f"[DEBUG forward_mixed] bsz={bsz}, total_q={total_q_tokens}, total_k={total_k_tokens}, expected_k={expected_k_len}" + ) + print(f"[DEBUG forward_mixed] has_prefix_cache={metadata.has_prefix_cache}") + print( + f"[DEBUG forward_mixed] cu_seqlens_q={forward_meta.cu_seqlens_q.tolist() if hasattr(forward_meta.cu_seqlens_q, 'tolist') else forward_meta.cu_seqlens_q}" + ) + if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: + print( + f"[DEBUG forward_mixed] cu_seqlens_k_with_cache={metadata.cu_seqlens_k_with_cache.tolist() if hasattr(metadata.cu_seqlens_k_with_cache, 'tolist') else metadata.cu_seqlens_k_with_cache}" + ) + + # Write cache only for new tokens prefill_mla_write_cache( compressed_kv, k_pe, @@ -581,22 +1019,31 @@ def forward_mixed( "none", self.max_seq_len, ) - # FA + + # Determine cu_seqlens_k and max_seqlen_k considering prefix cache + if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: + cu_seqlens_k = metadata.cu_seqlens_k_with_cache + max_seqlen_k = metadata.max_total_kv_len # max(dec_len + enc_len) + else: + cu_seqlens_k = forward_meta.cu_seqlens_k + max_seqlen_k = metadata.max_enc_len_this_time + + # FlashAttention for prefill fmha_out = self.flash_attn_func( q, k, v, forward_meta.cu_seqlens_q, - forward_meta.cu_seqlens_k, - metadata.max_enc_len_this_time, - metadata.max_enc_len_this_time, + cu_seqlens_k, + metadata.max_enc_len_this_time, # max_seqlen_q - only new tokens + max_seqlen_k, # max_seqlen_k - may include cached tokens causal=self.causal, **self.flash_attn_kwargs, )[0] return fmha_out - # Decode + # Decode branch: k is None if k is None: decode_mla_write_cache( compressed_kv, diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index 0b11b0b0905..0413d399279 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -351,7 +351,13 @@ def forward( forward_meta: ForwardMeta, hidden_states: paddle.Tensor, ): - """ """ + """MLA attention forward with prefix cache support.""" + + from fastdeploy.model_executor.layers.attention.mla_attention_backend import ( + MLAAttentionMetadata, + interleave_cached_and_new_latent, + read_latent_from_cache, + ) attn_out = None if self.use_gated_attn: @@ -378,7 +384,61 @@ def forward( need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0 if need_do_prefill: # max_enc_len_this_time - key_value = self.kv_b_proj(compressed_kv) + # Check for prefix cache + attn_meta = forward_meta.attn_backend.attention_metadata if hasattr(forward_meta, "attn_backend") else None + has_prefix_cache = False + total_cached_tokens = 0 + + if attn_meta is not None and isinstance(attn_meta, MLAAttentionMetadata): + has_prefix_cache = attn_meta.has_prefix_cache + total_cached_tokens = attn_meta.total_cached_kv_tokens + + # Handle prefix cache: read and concatenate cached latent with new latent + if has_prefix_cache and total_cached_tokens > 0: + # Get latent cache from forward_meta + layer_id = self.mla_attn.layer_id if hasattr(self.mla_attn, "layer_id") else 0 + latent_cache = forward_meta.caches[layer_id] if hasattr(forward_meta, "caches") else None + + if latent_cache is not None: + # Read cached latent vectors + cached_kv_c, cached_k_pe = read_latent_from_cache( + latent_cache, + forward_meta.block_tables, + forward_meta.seq_lens_decoder, # cache_kv_lens + attn_meta.cu_seqlens_cached_kv, + total_cached_tokens, + self.mla_attn.block_size if hasattr(self.mla_attn, "block_size") else 64, + self.kv_lora_rank, + self.qk_rope_head_dim, + ) + + if cached_kv_c is not None and cached_k_pe is not None: + # Interleave cached and new latent vectors per batch + # This ensures correct ordering: [batch0_cached, batch0_new, batch1_cached, batch1_new, ...] + # cached_k_pe already has RoPE applied (stored with RoPE) + full_compressed_kv, full_k_pe = interleave_cached_and_new_latent( + cached_kv_c, + cached_k_pe, + compressed_kv, + key_pe.squeeze(1), + attn_meta.cu_seqlens_cached_kv, + forward_meta.seq_lens_encoder, + forward_meta.cu_seqlens_q, + self.kv_lora_rank, + self.qk_rope_head_dim, + ) + else: + full_compressed_kv = compressed_kv + full_k_pe = key_pe.squeeze(1) + else: + full_compressed_kv = compressed_kv + full_k_pe = key_pe.squeeze(1) + else: + full_compressed_kv = compressed_kv + full_k_pe = key_pe.squeeze(1) + + # Project latent KV to full key and value + key_value = self.kv_b_proj(full_compressed_kv) key_value.reshape_( [ -1, @@ -389,9 +449,9 @@ def forward( key_nope, value = key_value.split([self.qk_nope_head_dim, self.v_head_dim], axis=-1) query[..., self.qk_nope_head_dim :] = query_pe - key = paddle.empty_like(query) + key = paddle.empty([full_k_pe.shape[0], self.num_attention_heads_tp, self.qk_head_dim], dtype=query.dtype) key[..., : self.qk_nope_head_dim] = key_nope - key[..., self.qk_nope_head_dim :] = key_pe + key[..., self.qk_nope_head_dim :] = full_k_pe.unsqueeze(1) value = paddle.nn.functional.pad(value, [0, self.qk_head_dim - self.v_head_dim], value=0) fmha_out = self.mla_attn( @@ -399,11 +459,15 @@ def forward( k=key, v=value, qkv=None, - compressed_kv=compressed_kv, - k_pe=key_pe, + compressed_kv=compressed_kv, # Pass original (new only) for cache writing + k_pe=key_pe, # Pass original (new only) for cache writing forward_meta=forward_meta, ) + # Debug: Print tensor shapes + print(f"[DEBUG deepseek_v3 forward] key.shape={key.shape}, value.shape={value.shape}") + print(f"[DEBUG deepseek_v3 forward] full_k_pe.shape={full_k_pe.shape if 'full_k_pe' in dir() else 'N/A'}") + fmha_out.reshape_([-1, self.num_attention_heads_tp, self.qk_head_dim]) fmha_out = fmha_out[:, :, : self.v_head_dim] fmha_out.reshape_([-1, self.num_attention_heads_tp * self.v_head_dim]) From c346a247e33a5e8fb85712f6aab97bffffe1dbc9 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Thu, 7 May 2026 14:22:11 +0800 Subject: [PATCH 02/10] optimized mla chunk prefill --- .../layers/attention/mla_attention_backend.py | 353 ++++++++++++++++-- .../model_executor/models/deepseek_v3.py | 8 +- 2 files changed, 324 insertions(+), 37 deletions(-) diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 153b2290250..3bbd99b4fc9 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -34,6 +34,12 @@ logger.debug(f"flash_attention_v3_varlen not available: {e}") flash_attention_v3_varlen = None +# Enable verbose debug logging for MLA prefix-cache / chunk-prefill paths when MLA_CHUNK_DEBUG=1. +# All logger.debug messages in this file are silent by default. +if os.environ.get("MLA_CHUNK_DEBUG", "0") == "1": + # paddleformers logger exposes set_level (not the standard logging.Logger.setLevel) + logger.set_level("DEBUG") + from fastdeploy.model_executor.layers.attention.ops import ( get_block_shape_and_split_kv_block, init_kv_signal_per_query, @@ -141,7 +147,7 @@ def read_latent_from_cache_kernel( tl.store(output_k_pe_base + k_pe_out_offsets, k_pe_value) -def read_latent_from_cache( +def read_latent_from_cache_naive( latent_cache: paddle.Tensor, block_tables: paddle.Tensor, cache_kv_lens: paddle.Tensor, @@ -182,9 +188,9 @@ def read_latent_from_cache( bsz = cu_seqlens_cached_kv.shape[0] - 1 output_idx = 0 - print(f"[DEBUG read_latent_from_cache] total_cached_tokens={total_cached_tokens}, bsz={bsz}") - print(f"[DEBUG read_latent_from_cache] cu_seqlens_cached_kv={cu_seqlens_cached_kv.tolist()}") - print(f"[DEBUG read_latent_from_cache] block_tables shape={block_tables.shape}") + logger.debug(f"[read_latent_from_cache] total_cached_tokens={total_cached_tokens}, bsz={bsz}") + logger.debug(f"[read_latent_from_cache] cu_seqlens_cached_kv={cu_seqlens_cached_kv.tolist()}") + logger.debug(f"[read_latent_from_cache] block_tables shape={block_tables.shape}") for batch_id in range(bsz): # Get the number of cached tokens for this batch from cu_seqlens_cached_kv @@ -204,7 +210,7 @@ def read_latent_from_cache( continue # Debug: Print cache reading info - print(f"[DEBUG read_latent_from_cache] batch_id={batch_id}, cache_len={cache_len}") + logger.debug(f"[read_latent_from_cache] batch_id={batch_id}, cache_len={cache_len}") # Read tokens from multiple blocks if cache_len > block_size local_idx = 0 @@ -216,8 +222,8 @@ def read_latent_from_cache( physical_block_id = block_tables[batch_id, block_idx].item() # Debug: Print block access info - print( - f"[DEBUG read_latent_from_cache] block_idx={block_idx}, block_offset={block_offset}, physical_block_id={physical_block_id}" + logger.debug( + f"[read_latent_from_cache] block_idx={block_idx}, block_offset={block_offset}, physical_block_id={physical_block_id}" ) # Load latent vectors from this block @@ -231,11 +237,14 @@ def read_latent_from_cache( local_idx += tokens_to_read - print(f"[DEBUG read_latent_from_cache] Total cached tokens read: {output_idx}") + logger.debug(f"[read_latent_from_cache] Total cached tokens read: {output_idx}") + assert ( + output_idx == total_cached_tokens + ), f"read_latent_from_cache_naive: wrote {output_idx} tokens, expected {total_cached_tokens}" return cached_kv_c, cached_k_pe -def interleave_cached_and_new_latent( +def interleave_cached_and_new_latent_naive( cached_kv_c: paddle.Tensor, cached_k_pe: paddle.Tensor, new_compressed_kv: paddle.Tensor, @@ -284,8 +293,8 @@ def interleave_cached_and_new_latent( new_idx = 0 out_position = 0 # Track output position for each batch - print( - f"[DEBUG interleave_cached_and_new_latent] bsz={bsz}, total_cached={total_cached}, total_new={total_new}, total_tokens={total_tokens}" + logger.debug( + f"[interleave_cached_and_new_latent] bsz={bsz}, total_cached={total_cached}, total_new={total_new}, total_tokens={total_tokens}" ) for batch_id in range(bsz): @@ -313,8 +322,8 @@ def interleave_cached_and_new_latent( ) num_new = cu_new_end - cu_new_start - print( - f"[DEBUG interleave] batch_id={batch_id}, num_cached={num_cached}, num_new={num_new}, cached_idx={cached_idx}, out_position={out_position}" + logger.debug( + f"[interleave] batch_id={batch_id}, num_cached={num_cached}, num_new={num_new}, cached_idx={cached_idx}, out_position={out_position}" ) # Output position for this batch (sequential, no gaps) @@ -339,10 +348,249 @@ def interleave_cached_and_new_latent( # Update output position for next batch out_position += num_cached + num_new - print(f"[DEBUG interleave] Final: cached_idx={cached_idx}, new_idx={new_idx}, out_position={out_position}") + logger.debug(f"[interleave] Final: cached_idx={cached_idx}, new_idx={new_idx}, out_position={out_position}") + assert ( + cached_idx == total_cached + ), f"interleave_cached_and_new_latent_naive: cached_idx={cached_idx} != total_cached={total_cached}" + assert new_idx == total_new, f"interleave_cached_and_new_latent_naive: new_idx={new_idx} != total_new={total_new}" + assert ( + out_position == total_tokens + ), f"interleave_cached_and_new_latent_naive: out_position={out_position} != total_tokens={total_tokens}" + return full_compressed_kv, full_k_pe + + +# ---------------------------------------------------------------------------- +# Public dispatchers. Default to the naive Python implementations; Task 3/4 will +# swap these to high-performance Triton kernels, controlled by FD_MLA_USE_NAIVE. +# ---------------------------------------------------------------------------- + + +# ---------------------------------------------------------------------------- +# Triton implementation of read_latent_from_cache. +# ---------------------------------------------------------------------------- + + +@enable_compat_on_triton_kernel +@triton.jit() +def _read_latent_triton_kernel( + latent_cache_ptr, # [num_blocks, 1, block_size, latent_dim] + block_tables_ptr, # [bsz, max_blocks_per_seq] + batch_id_per_token_ptr, # [total_cached_tokens] int32 + local_offset_per_token_ptr, # [total_cached_tokens] int32 + output_kv_c_ptr, # [total_cached_tokens, kv_lora_rank] + output_k_pe_ptr, # [total_cached_tokens, qk_rope_head_dim] + max_blocks_per_seq: tl.constexpr, + block_size: tl.constexpr, + kv_lora_rank: tl.constexpr, + qk_rope_head_dim: tl.constexpr, + LATENT_DIM: tl.constexpr, +): + token_idx = tl.program_id(axis=0) + + batch_id = tl.load(batch_id_per_token_ptr + token_idx) + local_off = tl.load(local_offset_per_token_ptr + token_idx) + + block_idx = local_off // block_size + block_offset = local_off % block_size + + physical_block_id = tl.load(block_tables_ptr + batch_id * max_blocks_per_seq + block_idx) + + latent_base = latent_cache_ptr + physical_block_id * block_size * LATENT_DIM + block_offset * LATENT_DIM + + kv_c_offs = tl.arange(0, kv_lora_rank) + kv_c_val = tl.load(latent_base + kv_c_offs) + tl.store(output_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) + + k_pe_offs = tl.arange(0, qk_rope_head_dim) + k_pe_val = tl.load(latent_base + kv_lora_rank + k_pe_offs) + tl.store(output_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) + + +def read_latent_from_cache_triton( + latent_cache: paddle.Tensor, + block_tables: paddle.Tensor, + cache_kv_lens: paddle.Tensor, + cu_seqlens_cached_kv: paddle.Tensor, + total_cached_tokens: int, + block_size: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +): + """Triton-accelerated version of read_latent_from_cache_naive. + + Returns (cached_kv_c, cached_k_pe) with the same semantics as the naive + implementation. + """ + if total_cached_tokens == 0: + return None, None + + bsz = cu_seqlens_cached_kv.shape[0] - 1 + max_blocks_per_seq = block_tables.shape[1] + + # --- host-side build of batch_id_per_token / local_offset_per_token --- + # O(bsz) CPU work; much cheaper than the in-kernel linear scan used by the + # legacy read_latent_from_cache_kernel. + cu_list = cu_seqlens_cached_kv.tolist() if hasattr(cu_seqlens_cached_kv, "tolist") else list(cu_seqlens_cached_kv) + batch_id_host = [0] * total_cached_tokens + local_off_host = [0] * total_cached_tokens + for b in range(bsz): + start = int(cu_list[b]) + end = int(cu_list[b + 1]) + for t in range(start, end): + batch_id_host[t] = b + local_off_host[t] = t - start + batch_id_per_token = paddle.to_tensor(batch_id_host, dtype="int32", place=latent_cache.place) + local_offset_per_token = paddle.to_tensor(local_off_host, dtype="int32", place=latent_cache.place) + + cached_kv_c = paddle.empty([total_cached_tokens, kv_lora_rank], dtype=latent_cache.dtype) + cached_k_pe = paddle.empty([total_cached_tokens, qk_rope_head_dim], dtype=latent_cache.dtype) + + grid = (total_cached_tokens,) + _read_latent_triton_kernel[grid]( + latent_cache, + block_tables, + batch_id_per_token, + local_offset_per_token, + cached_kv_c, + cached_k_pe, + max_blocks_per_seq=max_blocks_per_seq, + block_size=block_size, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + LATENT_DIM=kv_lora_rank + qk_rope_head_dim, + ) + + assert ( + cached_kv_c.shape[0] == total_cached_tokens + ), f"read_latent_from_cache_triton: output shape mismatch {cached_kv_c.shape[0]} vs {total_cached_tokens}" + return cached_kv_c, cached_k_pe + + +def read_latent_from_cache(*args, **kwargs): + if os.environ.get("FD_MLA_USE_NAIVE", "0") == "1": + return read_latent_from_cache_naive(*args, **kwargs) + return read_latent_from_cache_triton(*args, **kwargs) + + +# ---------------------------------------------------------------------------- +# Triton implementation of interleave_cached_and_new_latent. +# ---------------------------------------------------------------------------- + + +@enable_compat_on_triton_kernel +@triton.jit() +def _interleave_latent_kernel( + cached_kv_c_ptr, # [total_cached, kv_lora_rank] (may be invalid when total_cached==0) + cached_k_pe_ptr, # [total_cached, qk_rope_head_dim] + new_kv_c_ptr, # [total_new, kv_lora_rank] + new_k_pe_ptr, # [total_new, qk_rope_head_dim] + src_is_cached_ptr, # [total_tokens] int32 (1 == cached, 0 == new) + src_idx_ptr, # [total_tokens] int32 (index within the chosen source tensor) + out_kv_c_ptr, # [total_tokens, kv_lora_rank] + out_k_pe_ptr, # [total_tokens, qk_rope_head_dim] + kv_lora_rank: tl.constexpr, + qk_rope_head_dim: tl.constexpr, +): + token_idx = tl.program_id(axis=0) + is_cached = tl.load(src_is_cached_ptr + token_idx) + src_idx = tl.load(src_idx_ptr + token_idx) + + kv_c_offs = tl.arange(0, kv_lora_rank) + k_pe_offs = tl.arange(0, qk_rope_head_dim) + + if is_cached != 0: + kv_c_val = tl.load(cached_kv_c_ptr + src_idx * kv_lora_rank + kv_c_offs) + k_pe_val = tl.load(cached_k_pe_ptr + src_idx * qk_rope_head_dim + k_pe_offs) + else: + kv_c_val = tl.load(new_kv_c_ptr + src_idx * kv_lora_rank + kv_c_offs) + k_pe_val = tl.load(new_k_pe_ptr + src_idx * qk_rope_head_dim + k_pe_offs) + + tl.store(out_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) + tl.store(out_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) + + +def interleave_cached_and_new_latent_triton( + cached_kv_c: paddle.Tensor, + cached_k_pe: paddle.Tensor, + new_compressed_kv: paddle.Tensor, + new_k_pe: paddle.Tensor, + cu_seqlens_cached_kv: paddle.Tensor, + seq_lens_encoder: paddle.Tensor, + cu_seqlens_q: paddle.Tensor, + kv_lora_rank: int, + qk_rope_head_dim: int, +): + """Triton-accelerated version of interleave_cached_and_new_latent_naive.""" + bsz = cu_seqlens_cached_kv.shape[0] - 1 + total_cached = cached_kv_c.shape[0] if cached_kv_c is not None and cached_kv_c.numel() > 0 else 0 + total_new = new_compressed_kv.shape[0] + total_tokens = total_cached + total_new + + full_compressed_kv = paddle.empty([total_tokens, kv_lora_rank], dtype=new_compressed_kv.dtype) + full_k_pe = paddle.empty([total_tokens, qk_rope_head_dim], dtype=new_k_pe.dtype) + + if total_tokens == 0: + return full_compressed_kv, full_k_pe + + # ---- host-side mapping: for each output position, pick source ---- + cu_cached = ( + cu_seqlens_cached_kv.tolist() if hasattr(cu_seqlens_cached_kv, "tolist") else list(cu_seqlens_cached_kv) + ) + cu_new = cu_seqlens_q.tolist() if hasattr(cu_seqlens_q, "tolist") else list(cu_seqlens_q) + + src_is_cached = [0] * total_tokens + src_idx = [0] * total_tokens + out_pos = 0 + for b in range(bsz): + nc = int(cu_cached[b + 1]) - int(cu_cached[b]) + nn = int(cu_new[b + 1]) - int(cu_new[b]) + cached_base = int(cu_cached[b]) + new_base = int(cu_new[b]) + for t in range(nc): + src_is_cached[out_pos] = 1 + src_idx[out_pos] = cached_base + t + out_pos += 1 + for t in range(nn): + src_is_cached[out_pos] = 0 + src_idx[out_pos] = new_base + t + out_pos += 1 + + assert ( + out_pos == total_tokens + ), f"interleave_cached_and_new_latent_triton: host map out_pos={out_pos} != total_tokens={total_tokens}" + + dev = new_compressed_kv.place + src_is_cached_t = paddle.to_tensor(src_is_cached, dtype="int32", place=dev) + src_idx_t = paddle.to_tensor(src_idx, dtype="int32", place=dev) + + # Provide non-null pointers for the cached tensors even when total_cached == 0; + # Triton kernel path is still gated by src_is_cached flag. + cached_kv_c_ptr = cached_kv_c if total_cached > 0 else full_compressed_kv + cached_k_pe_ptr = cached_k_pe if total_cached > 0 else full_k_pe + + grid = (total_tokens,) + _interleave_latent_kernel[grid]( + cached_kv_c_ptr, + cached_k_pe_ptr, + new_compressed_kv, + new_k_pe, + src_is_cached_t, + src_idx_t, + full_compressed_kv, + full_k_pe, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + ) + return full_compressed_kv, full_k_pe +def interleave_cached_and_new_latent(*args, **kwargs): + if os.environ.get("FD_MLA_USE_NAIVE", "0") == "1": + return interleave_cached_and_new_latent_naive(*args, **kwargs) + return interleave_cached_and_new_latent_triton(*args, **kwargs) + + # ============================================================================ @@ -680,18 +928,22 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): if hasattr(forward_meta.seq_lens_decoder[i], "item") else forward_meta.seq_lens_decoder[i] ) - # Count cached tokens whenever dec_len > 0, which includes: - # - chunked-prefill continuation (enc_len>0, dec_len>0) - # - pure decode (enc_len==0, dec_len>0) — its KV is read via this path too so that - # cu_seqlens_k_with_cache stays consistent with cu_seqlens_q for the prefill - # FlashAttention call; decode batch outputs here are discarded by merge_prefill_decode_output. + seq_this_time = ( + forward_meta.seq_lens_this_time[i].item() + if hasattr(forward_meta.seq_lens_this_time[i], "item") + else forward_meta.seq_lens_this_time[i] + ) + # Per-batch K length after interleave = dec_len (cached) + seq_this_time (new). + # seq_this_time equals enc_len for prefill/chunked-prefill, and 1 for decode. + # Using enc_len here drops the decode-batch new token and yields a wrong + # max_seqlen_k for FlashAttention, which then tile-truncates the last K row. + per_batch_k = dec_len + seq_this_time if dec_len > 0: has_prefix_cache = True total_cached_kv_tokens += dec_len max_cached_kv_len = max(max_cached_kv_len, dec_len) - max_total_kv_len = max(max_total_kv_len, dec_len + enc_len) - if enc_len > 0: - max_total_kv_len = max(max_total_kv_len, max(enc_len, dec_len + enc_len)) + if per_batch_k > 0: + max_total_kv_len = max(max_total_kv_len, per_batch_k) metadata.has_prefix_cache = has_prefix_cache metadata.total_cached_kv_tokens = total_cached_kv_tokens @@ -729,8 +981,8 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): if hasattr(forward_meta.seq_lens_this_time[i], "item") else forward_meta.seq_lens_this_time[i] ) - print( - f"[DEBUG init_attn_meta] batch {i}: enc_len={enc_len}, dec_len={dec_len}, seq_this={seq_this_time}, cumsum_cached={cumsum_cached}, cumsum_total={cumsum_total}" + logger.debug( + f"[init_attn_meta] batch {i}: enc_len={enc_len}, dec_len={dec_len}, seq_this={seq_this_time}, cumsum_cached={cumsum_cached}, cumsum_total={cumsum_total}" ) if dec_len > 0: cumsum_cached += dec_len @@ -740,10 +992,30 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): cumsum_total += seq_this_time cu_seqlens_cached_kv[i + 1] = cumsum_cached cu_seqlens_k_with_cache[i + 1] = cumsum_total - print(f"[DEBUG init_attn_meta] Final cu_seqlens_cached_kv: {cu_seqlens_cached_kv.tolist()}") - print(f"[DEBUG init_attn_meta] Final cu_seqlens_k_with_cache: {cu_seqlens_k_with_cache.tolist()}") + logger.debug(f"[init_attn_meta] Final cu_seqlens_cached_kv: {cu_seqlens_cached_kv.tolist()}") + logger.debug(f"[init_attn_meta] Final cu_seqlens_k_with_cache: {cu_seqlens_k_with_cache.tolist()}") + # Consistency checks: starts at 0, monotonic non-decreasing, final equals cumulative. + assert cu_seqlens_cached_kv[0].item() == 0, "cu_seqlens_cached_kv must start at 0" + assert cu_seqlens_k_with_cache[0].item() == 0, "cu_seqlens_k_with_cache must start at 0" + assert ( + cu_seqlens_cached_kv[bsz].item() == cumsum_cached + ), f"cu_seqlens_cached_kv[-1]={cu_seqlens_cached_kv[bsz].item()} != cumsum_cached={cumsum_cached}" + assert ( + cu_seqlens_k_with_cache[bsz].item() == cumsum_total + ), f"cu_seqlens_k_with_cache[-1]={cu_seqlens_k_with_cache[bsz].item()} != cumsum_total={cumsum_total}" metadata.cu_seqlens_cached_kv = cu_seqlens_cached_kv metadata.cu_seqlens_k_with_cache = cu_seqlens_k_with_cache + # Cross-check: max_total_kv_len MUST be >= max per-batch K length implied by + # cu_seqlens_k_with_cache. Otherwise FlashAttention will tile-truncate the + # last K rows of the longest batch and silently corrupt attention output. + if bsz > 0: + cu_k_list = cu_seqlens_k_with_cache.tolist() + observed_max = max(cu_k_list[i + 1] - cu_k_list[i] for i in range(bsz)) + assert max_total_kv_len >= observed_max, ( + f"max_total_kv_len={max_total_kv_len} < observed per-batch max " + f"K length={observed_max} from cu_seqlens_k_with_cache={cu_k_list}. " + f"FlashAttention will truncate K tiles." + ) # pd_disaggregation metadata.kv_signal_data_list = [None] * self.num_layers @@ -993,19 +1265,23 @@ def forward_mixed( ) # Debug output - print( - f"[DEBUG forward_mixed] bsz={bsz}, total_q={total_q_tokens}, total_k={total_k_tokens}, expected_k={expected_k_len}" + logger.debug( + f"[forward_mixed] bsz={bsz}, total_q={total_q_tokens}, total_k={total_k_tokens}, expected_k={expected_k_len}" ) - print(f"[DEBUG forward_mixed] has_prefix_cache={metadata.has_prefix_cache}") - print( - f"[DEBUG forward_mixed] cu_seqlens_q={forward_meta.cu_seqlens_q.tolist() if hasattr(forward_meta.cu_seqlens_q, 'tolist') else forward_meta.cu_seqlens_q}" + logger.debug(f"[forward_mixed] has_prefix_cache={metadata.has_prefix_cache}") + logger.debug( + f"[forward_mixed] cu_seqlens_q={forward_meta.cu_seqlens_q.tolist() if hasattr(forward_meta.cu_seqlens_q, 'tolist') else forward_meta.cu_seqlens_q}" ) if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: - print( - f"[DEBUG forward_mixed] cu_seqlens_k_with_cache={metadata.cu_seqlens_k_with_cache.tolist() if hasattr(metadata.cu_seqlens_k_with_cache, 'tolist') else metadata.cu_seqlens_k_with_cache}" + logger.debug( + f"[forward_mixed] cu_seqlens_k_with_cache={metadata.cu_seqlens_k_with_cache.tolist() if hasattr(metadata.cu_seqlens_k_with_cache, 'tolist') else metadata.cu_seqlens_k_with_cache}" ) - # Write cache only for new tokens + # Write cache only for new tokens of prefill/chunked-prefill batches. + # Decode batches (seq_lens_encoder == 0) are intentionally skipped here — they + # are handled later by decode_mla_write_cache() via the k=None branch when + # need_do_decode is True. Using seq_lens_encoder keeps that separation correct + # and avoids double-writing decode tokens. prefill_mla_write_cache( compressed_kv, k_pe, @@ -1028,6 +1304,15 @@ def forward_mixed( cu_seqlens_k = forward_meta.cu_seqlens_k max_seqlen_k = metadata.max_enc_len_this_time + # Shape consistency: k/v token count must match cu_seqlens_k terminal value + _expected_k = cu_seqlens_k[bsz].item() if hasattr(cu_seqlens_k[bsz], "item") else int(cu_seqlens_k[bsz]) + assert ( + k.shape[0] == _expected_k + ), f"forward_mixed: k.shape[0]={k.shape[0]} != cu_seqlens_k[-1]={_expected_k}" + assert ( + v.shape[0] == _expected_k + ), f"forward_mixed: v.shape[0]={v.shape[0]} != cu_seqlens_k[-1]={_expected_k}" + # FlashAttention for prefill fmha_out = self.flash_attn_func( q, diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index 0413d399279..2261ab34c42 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -464,9 +464,11 @@ def forward( forward_meta=forward_meta, ) - # Debug: Print tensor shapes - print(f"[DEBUG deepseek_v3 forward] key.shape={key.shape}, value.shape={value.shape}") - print(f"[DEBUG deepseek_v3 forward] full_k_pe.shape={full_k_pe.shape if 'full_k_pe' in dir() else 'N/A'}") + # Gated by MLA_CHUNK_DEBUG=1 via logger.debug (see mla_attention_backend.py). + logger.debug( + f"[deepseek_v3 forward] key.shape={key.shape}, value.shape={value.shape}, " + f"full_k_pe.shape={full_k_pe.shape}" + ) fmha_out.reshape_([-1, self.num_attention_heads_tp, self.qk_head_dim]) fmha_out = fmha_out[:, :, : self.v_head_dim] From aaa5924fb5a8054d267abd30a409c35d4d1616d0 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Thu, 7 May 2026 15:43:07 +0800 Subject: [PATCH 03/10] optimized mla chunk prefill --- .../layers/attention/mla_attention_backend.py | 156 ++---------------- .../model_executor/models/deepseek_v3.py | 6 - 2 files changed, 12 insertions(+), 150 deletions(-) diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 3bbd99b4fc9..6150394a15e 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -34,12 +34,6 @@ logger.debug(f"flash_attention_v3_varlen not available: {e}") flash_attention_v3_varlen = None -# Enable verbose debug logging for MLA prefix-cache / chunk-prefill paths when MLA_CHUNK_DEBUG=1. -# All logger.debug messages in this file are silent by default. -if os.environ.get("MLA_CHUNK_DEBUG", "0") == "1": - # paddleformers logger exposes set_level (not the standard logging.Logger.setLevel) - logger.set_level("DEBUG") - from fastdeploy.model_executor.layers.attention.ops import ( get_block_shape_and_split_kv_block, init_kv_signal_per_query, @@ -78,75 +72,6 @@ # ============================================================================ -@enable_compat_on_triton_kernel -@triton.jit() -def read_latent_from_cache_kernel( - latent_cache, - block_tables, - cache_kv_lens, - cu_seqlens_cached_kv, - output_kv_c, - output_k_pe, - block_size: tl.constexpr, - kv_lora_rank: tl.constexpr, - qk_rope_head_dim: tl.constexpr, - LATENT_DIM: tl.constexpr, -): - """ - Kernel to read latent vectors (kv_c and k_pe) from paged latent cache. - Each program instance handles one cached token. - - Args: - latent_cache: [num_blocks, 1, block_size, kv_lora_rank + qk_rope_head_dim] - block_tables: [batch_size, max_blocks_per_seq] - cache_kv_lens: [batch_size] - cached KV length for each request - cu_seqlens_cached_kv: [batch_size + 1] - cumulative sequence lengths for cached KV - output_kv_c: [total_cached_tokens, kv_lora_rank] - output_k_pe: [total_cached_tokens, qk_rope_head_dim] - """ - # Global token index in the output - token_idx = tl.program_id(axis=0) - - # Find which batch this token belongs to using binary search on cu_seqlens - # For simplicity, we use a linear scan (could be optimized with binary search) - batch_id = 0 - for i in range(cu_seqlens_cached_kv.shape[0] - 1): - if token_idx >= tl.load(cu_seqlens_cached_kv + i) and token_idx < tl.load(cu_seqlens_cached_kv + i + 1): - batch_id = i - break - - # Local token index within the batch - cu_start = tl.load(cu_seqlens_cached_kv + batch_id) - local_token_idx = token_idx - cu_start - - # Get the physical block and offset - block_idx = local_token_idx // block_size - block_offset = local_token_idx % block_size - - # Get physical block id from block_tables - physical_block_id = tl.load(block_tables + batch_id * block_tables.shape[1] + block_idx) - - # Load latent vector from cache - # latent_cache shape: [num_blocks, 1, block_size, kv_lora_rank + qk_rope_head_dim] - latent_base = latent_cache + physical_block_id * LATENT_DIM * block_size + block_offset * LATENT_DIM - - # Read kv_c (first kv_lora_rank dimensions) - kv_c_offsets = tl.arange(0, kv_lora_rank) - kv_c_value = tl.load(latent_base + kv_c_offsets, mask=kv_c_offsets < kv_lora_rank) - - # Read k_pe (last qk_rope_head_dim dimensions) - k_pe_offsets = tl.arange(kv_lora_rank, kv_lora_rank + qk_rope_head_dim) - k_pe_value = tl.load(latent_base + k_pe_offsets, mask=k_pe_offsets < LATENT_DIM) - - # Store outputs - output_kv_c_base = output_kv_c + token_idx * kv_lora_rank - tl.store(output_kv_c_base + kv_c_offsets, kv_c_value, mask=kv_c_offsets < kv_lora_rank) - - output_k_pe_base = output_k_pe + token_idx * qk_rope_head_dim - k_pe_out_offsets = tl.arange(0, qk_rope_head_dim) - tl.store(output_k_pe_base + k_pe_out_offsets, k_pe_value) - - def read_latent_from_cache_naive( latent_cache: paddle.Tensor, block_tables: paddle.Tensor, @@ -188,10 +113,6 @@ def read_latent_from_cache_naive( bsz = cu_seqlens_cached_kv.shape[0] - 1 output_idx = 0 - logger.debug(f"[read_latent_from_cache] total_cached_tokens={total_cached_tokens}, bsz={bsz}") - logger.debug(f"[read_latent_from_cache] cu_seqlens_cached_kv={cu_seqlens_cached_kv.tolist()}") - logger.debug(f"[read_latent_from_cache] block_tables shape={block_tables.shape}") - for batch_id in range(bsz): # Get the number of cached tokens for this batch from cu_seqlens_cached_kv cu_start = ( @@ -209,9 +130,6 @@ def read_latent_from_cache_naive( if cache_len <= 0: continue - # Debug: Print cache reading info - logger.debug(f"[read_latent_from_cache] batch_id={batch_id}, cache_len={cache_len}") - # Read tokens from multiple blocks if cache_len > block_size local_idx = 0 while local_idx < cache_len: @@ -221,11 +139,6 @@ def read_latent_from_cache_naive( physical_block_id = block_tables[batch_id, block_idx].item() - # Debug: Print block access info - logger.debug( - f"[read_latent_from_cache] block_idx={block_idx}, block_offset={block_offset}, physical_block_id={physical_block_id}" - ) - # Load latent vectors from this block for offset in range(tokens_to_read): latent_vec = latent_cache[physical_block_id, 0, block_offset + offset, :] @@ -237,7 +150,6 @@ def read_latent_from_cache_naive( local_idx += tokens_to_read - logger.debug(f"[read_latent_from_cache] Total cached tokens read: {output_idx}") assert ( output_idx == total_cached_tokens ), f"read_latent_from_cache_naive: wrote {output_idx} tokens, expected {total_cached_tokens}" @@ -293,10 +205,6 @@ def interleave_cached_and_new_latent_naive( new_idx = 0 out_position = 0 # Track output position for each batch - logger.debug( - f"[interleave_cached_and_new_latent] bsz={bsz}, total_cached={total_cached}, total_new={total_new}, total_tokens={total_tokens}" - ) - for batch_id in range(bsz): # Number of cached tokens for this batch cu_cached_start = ( @@ -322,10 +230,6 @@ def interleave_cached_and_new_latent_naive( ) num_new = cu_new_end - cu_new_start - logger.debug( - f"[interleave] batch_id={batch_id}, num_cached={num_cached}, num_new={num_new}, cached_idx={cached_idx}, out_position={out_position}" - ) - # Output position for this batch (sequential, no gaps) out_start = out_position @@ -348,7 +252,6 @@ def interleave_cached_and_new_latent_naive( # Update output position for next batch out_position += num_cached + num_new - logger.debug(f"[interleave] Final: cached_idx={cached_idx}, new_idx={new_idx}, out_position={out_position}") assert ( cached_idx == total_cached ), f"interleave_cached_and_new_latent_naive: cached_idx={cached_idx} != total_cached={total_cached}" @@ -852,12 +755,12 @@ def __init__( is_paddle_supported = any(num >= 90 for num in paddle.version.cuda_archs()) if is_current_sm_supported and is_paddle_supported: self.flash_attn_func = flash_attention_v3_varlen - print("The current platform supports Flash Attention V3.") + logger.info("The current platform supports Flash Attention V3.") self.flash_attn_kwargs = {"softmax_scale": self.attn_softmax_scale} else: self.flash_attn_func = flash_attn_unpadded self.flash_attn_kwargs = {"scale": self.attn_softmax_scale, "training": False} - print( + logger.info( "The current platform does not support Flash Attention V3, so Flash Attention V2 will be used instead." ) @@ -918,11 +821,11 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): # Prefix cache exists when seq_lens_decoder > 0 # seq_lens_decoder stores the cached KV length for chunked prefill/prefix cache for i in range(bsz): - enc_len = ( - forward_meta.seq_lens_encoder[i].item() - if hasattr(forward_meta.seq_lens_encoder[i], "item") - else forward_meta.seq_lens_encoder[i] - ) + # enc_len = ( + # forward_meta.seq_lens_encoder[i].item() + # if hasattr(forward_meta.seq_lens_encoder[i], "item") + # else forward_meta.seq_lens_encoder[i] + # ) dec_len = ( forward_meta.seq_lens_decoder[i].item() if hasattr(forward_meta.seq_lens_decoder[i], "item") @@ -966,11 +869,11 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): # cu_seqlens_k_with_cache must reflect this sum per batch. # cu_seqlens_cached_kv tracks only the cached portion for read_latent_from_cache(). for i in range(bsz): - enc_len = ( - forward_meta.seq_lens_encoder[i].item() - if hasattr(forward_meta.seq_lens_encoder[i], "item") - else forward_meta.seq_lens_encoder[i] - ) + # enc_len = ( + # forward_meta.seq_lens_encoder[i].item() + # if hasattr(forward_meta.seq_lens_encoder[i], "item") + # else forward_meta.seq_lens_encoder[i] + # ) dec_len = ( forward_meta.seq_lens_decoder[i].item() if hasattr(forward_meta.seq_lens_decoder[i], "item") @@ -981,9 +884,6 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): if hasattr(forward_meta.seq_lens_this_time[i], "item") else forward_meta.seq_lens_this_time[i] ) - logger.debug( - f"[init_attn_meta] batch {i}: enc_len={enc_len}, dec_len={dec_len}, seq_this={seq_this_time}, cumsum_cached={cumsum_cached}, cumsum_total={cumsum_total}" - ) if dec_len > 0: cumsum_cached += dec_len cumsum_total += dec_len @@ -992,8 +892,6 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): cumsum_total += seq_this_time cu_seqlens_cached_kv[i + 1] = cumsum_cached cu_seqlens_k_with_cache[i + 1] = cumsum_total - logger.debug(f"[init_attn_meta] Final cu_seqlens_cached_kv: {cu_seqlens_cached_kv.tolist()}") - logger.debug(f"[init_attn_meta] Final cu_seqlens_k_with_cache: {cu_seqlens_k_with_cache.tolist()}") # Consistency checks: starts at 0, monotonic non-decreasing, final equals cumulative. assert cu_seqlens_cached_kv[0].item() == 0, "cu_seqlens_cached_kv must start at 0" assert cu_seqlens_k_with_cache[0].item() == 0, "cu_seqlens_k_with_cache must start at 0" @@ -1245,37 +1143,7 @@ def forward_mixed( # Prefill branch: k is not None if k is not None: - # Debug: Verify tensor shapes and sequence lengths bsz = forward_meta.cu_seqlens_q.shape[0] - 1 - total_q_tokens = q.shape[0] - total_k_tokens = k.shape[0] - - # Calculate expected cu_seqlens_k_with_cache - if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: - expected_k_len = ( - metadata.cu_seqlens_k_with_cache[bsz].item() - if hasattr(metadata.cu_seqlens_k_with_cache[bsz], "item") - else metadata.cu_seqlens_k_with_cache[bsz] - ) - else: - expected_k_len = ( - forward_meta.cu_seqlens_k[bsz].item() - if hasattr(forward_meta.cu_seqlens_k[bsz], "item") - else forward_meta.cu_seqlens_k[bsz] - ) - - # Debug output - logger.debug( - f"[forward_mixed] bsz={bsz}, total_q={total_q_tokens}, total_k={total_k_tokens}, expected_k={expected_k_len}" - ) - logger.debug(f"[forward_mixed] has_prefix_cache={metadata.has_prefix_cache}") - logger.debug( - f"[forward_mixed] cu_seqlens_q={forward_meta.cu_seqlens_q.tolist() if hasattr(forward_meta.cu_seqlens_q, 'tolist') else forward_meta.cu_seqlens_q}" - ) - if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: - logger.debug( - f"[forward_mixed] cu_seqlens_k_with_cache={metadata.cu_seqlens_k_with_cache.tolist() if hasattr(metadata.cu_seqlens_k_with_cache, 'tolist') else metadata.cu_seqlens_k_with_cache}" - ) # Write cache only for new tokens of prefill/chunked-prefill batches. # Decode batches (seq_lens_encoder == 0) are intentionally skipped here — they diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index 2261ab34c42..78d05aabcb1 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -464,12 +464,6 @@ def forward( forward_meta=forward_meta, ) - # Gated by MLA_CHUNK_DEBUG=1 via logger.debug (see mla_attention_backend.py). - logger.debug( - f"[deepseek_v3 forward] key.shape={key.shape}, value.shape={value.shape}, " - f"full_k_pe.shape={full_k_pe.shape}" - ) - fmha_out.reshape_([-1, self.num_attention_heads_tp, self.qk_head_dim]) fmha_out = fmha_out[:, :, : self.v_head_dim] fmha_out.reshape_([-1, self.num_attention_heads_tp * self.v_head_dim]) From 03ac41e82c7541c87550a962ed540f98f7cc9f1e Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Thu, 7 May 2026 17:17:18 +0800 Subject: [PATCH 04/10] update mla chunk-prefill --- .../layers/attention/mla_attention_backend.py | 432 ++++-------------- .../model_executor/models/deepseek_v3.py | 46 +- 2 files changed, 103 insertions(+), 375 deletions(-) diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 6150394a15e..050e3e90119 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -68,433 +68,184 @@ from fastdeploy.spec_decode import SpecMethod # ============================================================================ -# Latent Cache Read Kernel for Prefix Cache Support +# Fused Read-Cache + Interleave Kernel for Prefix Cache Support +# +# For each output token position t in [0, total_tokens): +# - if cached: load latent vector from paged latent_cache (physical block) +# - if new: load from new_compressed_kv / new_k_pe at the given index +# A single kernel produces full_compressed_kv / full_k_pe directly, avoiding +# an intermediate (cached_kv_c, cached_k_pe) allocation and an extra launch. # ============================================================================ -def read_latent_from_cache_naive( +def fused_read_cache_and_interleave_naive( latent_cache: paddle.Tensor, block_tables: paddle.Tensor, - cache_kv_lens: paddle.Tensor, - cu_seqlens_cached_kv: paddle.Tensor, - total_cached_tokens: int, - block_size: int, - kv_lora_rank: int, - qk_rope_head_dim: int, -) -> Tuple[paddle.Tensor, paddle.Tensor]: - """ - Read latent vectors (kv_c and k_pe) from paged latent cache for prefix cache support. - - Args: - latent_cache: [num_blocks, 1, block_size, kv_lora_rank + qk_rope_head_dim] - block_tables: [batch_size, max_blocks_per_seq] - cache_kv_lens: [batch_size] - cached KV length for each request - cu_seqlens_cached_kv: [batch_size + 1] - cumulative sequence lengths for cached KV - total_cached_tokens: Total number of cached tokens across all requests - block_size: Block size for paged attention - kv_lora_rank: LoRA rank for KV compression - qk_rope_head_dim: Dimension of RoPE part in key - - Returns: - cached_kv_c: [total_cached_tokens, kv_lora_rank] - cached_k_pe: [total_cached_tokens, qk_rope_head_dim] - """ - if total_cached_tokens == 0: - return None, None - - # latent_dim = kv_lora_rank + qk_rope_head_dim - - # Allocate output tensors - cached_kv_c = paddle.empty([total_cached_tokens, kv_lora_rank], dtype=latent_cache.dtype) - cached_k_pe = paddle.empty([total_cached_tokens, qk_rope_head_dim], dtype=latent_cache.dtype) - - # Use a simpler approach: iterate through each batch and gather latent vectors - # IMPORTANT: Only process batches that have prefix cache (use cu_seqlens_cached_kv) - bsz = cu_seqlens_cached_kv.shape[0] - 1 - output_idx = 0 - - for batch_id in range(bsz): - # Get the number of cached tokens for this batch from cu_seqlens_cached_kv - cu_start = ( - cu_seqlens_cached_kv[batch_id].item() - if hasattr(cu_seqlens_cached_kv[batch_id], "item") - else cu_seqlens_cached_kv[batch_id] - ) - cu_end = ( - cu_seqlens_cached_kv[batch_id + 1].item() - if hasattr(cu_seqlens_cached_kv[batch_id + 1], "item") - else cu_seqlens_cached_kv[batch_id + 1] - ) - cache_len = cu_end - cu_start - - if cache_len <= 0: - continue - - # Read tokens from multiple blocks if cache_len > block_size - local_idx = 0 - while local_idx < cache_len: - block_idx = local_idx // block_size - block_offset = local_idx % block_size - tokens_to_read = min(block_size - block_offset, cache_len - local_idx) - - physical_block_id = block_tables[batch_id, block_idx].item() - - # Load latent vectors from this block - for offset in range(tokens_to_read): - latent_vec = latent_cache[physical_block_id, 0, block_offset + offset, :] - - # Split into kv_c and k_pe - cached_kv_c[output_idx] = latent_vec[:kv_lora_rank] - cached_k_pe[output_idx] = latent_vec[kv_lora_rank:] - output_idx += 1 - - local_idx += tokens_to_read - - assert ( - output_idx == total_cached_tokens - ), f"read_latent_from_cache_naive: wrote {output_idx} tokens, expected {total_cached_tokens}" - return cached_kv_c, cached_k_pe - - -def interleave_cached_and_new_latent_naive( - cached_kv_c: paddle.Tensor, - cached_k_pe: paddle.Tensor, new_compressed_kv: paddle.Tensor, new_k_pe: paddle.Tensor, cu_seqlens_cached_kv: paddle.Tensor, - seq_lens_encoder: paddle.Tensor, cu_seqlens_q: paddle.Tensor, kv_lora_rank: int, qk_rope_head_dim: int, + block_size: int, ) -> Tuple[paddle.Tensor, paddle.Tensor]: - """ - Interleave cached and new latent vectors per batch for correct FlashAttention layout. - - The key insight is that FlashAttention expects tokens to be ordered by batch: - [batch0_tokens, batch1_tokens, ...] - where each batch's tokens should be [cached_tokens, new_tokens]. - - Args: - cached_kv_c: [total_cached_tokens, kv_lora_rank] - cached KV latent (from read_latent_from_cache) - cached_k_pe: [total_cached_tokens, qk_rope_head_dim] - cached K RoPE (from read_latent_from_cache) - new_compressed_kv: [total_new_tokens, kv_lora_rank] - new KV latent from current forward - new_k_pe: [total_new_tokens, qk_rope_head_dim] - new K RoPE from current forward (already with RoPE applied) - cu_seqlens_cached_kv: [batch_size + 1] - cumulative sequence lengths for cached KV - seq_lens_encoder: [batch_size] - number of new tokens per batch - cu_seqlens_q: [batch_size + 1] - cumulative sequence lengths for new tokens - kv_lora_rank: LoRA rank for KV compression - qk_rope_head_dim: Dimension of RoPE part in key - - Returns: - full_compressed_kv: [total_tokens, kv_lora_rank] - properly interleaved KV latent - full_k_pe: [total_tokens, qk_rope_head_dim] - properly interleaved K RoPE + """Python/Paddle reference of the fused read+interleave path. + + Returns ``(full_compressed_kv, full_k_pe)`` with per-batch layout + ``[cached_tokens, new_tokens]`` in token order. """ bsz = cu_seqlens_cached_kv.shape[0] - 1 - - # Calculate total output size - total_cached = cached_kv_c.shape[0] if cached_kv_c is not None and cached_kv_c.numel() > 0 else 0 + cu_cached = cu_seqlens_cached_kv.tolist() + cu_new = cu_seqlens_q.tolist() + total_cached = int(cu_cached[bsz]) total_new = new_compressed_kv.shape[0] total_tokens = total_cached + total_new - # Allocate output tensors full_compressed_kv = paddle.empty([total_tokens, kv_lora_rank], dtype=new_compressed_kv.dtype) full_k_pe = paddle.empty([total_tokens, qk_rope_head_dim], dtype=new_k_pe.dtype) + if total_tokens == 0: + return full_compressed_kv, full_k_pe - # Track indices into cached and new tensors - cached_idx = 0 - new_idx = 0 - out_position = 0 # Track output position for each batch - - for batch_id in range(bsz): - # Number of cached tokens for this batch - cu_cached_start = ( - cu_seqlens_cached_kv[batch_id].item() - if hasattr(cu_seqlens_cached_kv[batch_id], "item") - else cu_seqlens_cached_kv[batch_id] - ) - cu_cached_end = ( - cu_seqlens_cached_kv[batch_id + 1].item() - if hasattr(cu_seqlens_cached_kv[batch_id + 1], "item") - else cu_seqlens_cached_kv[batch_id + 1] - ) - num_cached = cu_cached_end - cu_cached_start - - # Number of new tokens for this batch - cu_new_start = ( - cu_seqlens_q[batch_id].item() if hasattr(cu_seqlens_q[batch_id], "item") else cu_seqlens_q[batch_id] - ) - cu_new_end = ( - cu_seqlens_q[batch_id + 1].item() - if hasattr(cu_seqlens_q[batch_id + 1], "item") - else cu_seqlens_q[batch_id + 1] - ) - num_new = cu_new_end - cu_new_start - - # Output position for this batch (sequential, no gaps) - out_start = out_position - - # Copy cached tokens first (if any) - if num_cached > 0 and cached_kv_c is not None: - full_compressed_kv[out_start : out_start + num_cached] = cached_kv_c[cached_idx : cached_idx + num_cached] - full_k_pe[out_start : out_start + num_cached] = cached_k_pe[cached_idx : cached_idx + num_cached] - cached_idx += num_cached - - # Then copy new tokens - if num_new > 0: - full_compressed_kv[out_start + num_cached : out_start + num_cached + num_new] = new_compressed_kv[ - new_idx : new_idx + num_new - ] - full_k_pe[out_start + num_cached : out_start + num_cached + num_new] = new_k_pe[ - new_idx : new_idx + num_new - ] - new_idx += num_new - - # Update output position for next batch - out_position += num_cached + num_new - - assert ( - cached_idx == total_cached - ), f"interleave_cached_and_new_latent_naive: cached_idx={cached_idx} != total_cached={total_cached}" - assert new_idx == total_new, f"interleave_cached_and_new_latent_naive: new_idx={new_idx} != total_new={total_new}" - assert ( - out_position == total_tokens - ), f"interleave_cached_and_new_latent_naive: out_position={out_position} != total_tokens={total_tokens}" - return full_compressed_kv, full_k_pe - - -# ---------------------------------------------------------------------------- -# Public dispatchers. Default to the naive Python implementations; Task 3/4 will -# swap these to high-performance Triton kernels, controlled by FD_MLA_USE_NAIVE. -# ---------------------------------------------------------------------------- - - -# ---------------------------------------------------------------------------- -# Triton implementation of read_latent_from_cache. -# ---------------------------------------------------------------------------- - - -@enable_compat_on_triton_kernel -@triton.jit() -def _read_latent_triton_kernel( - latent_cache_ptr, # [num_blocks, 1, block_size, latent_dim] - block_tables_ptr, # [bsz, max_blocks_per_seq] - batch_id_per_token_ptr, # [total_cached_tokens] int32 - local_offset_per_token_ptr, # [total_cached_tokens] int32 - output_kv_c_ptr, # [total_cached_tokens, kv_lora_rank] - output_k_pe_ptr, # [total_cached_tokens, qk_rope_head_dim] - max_blocks_per_seq: tl.constexpr, - block_size: tl.constexpr, - kv_lora_rank: tl.constexpr, - qk_rope_head_dim: tl.constexpr, - LATENT_DIM: tl.constexpr, -): - token_idx = tl.program_id(axis=0) - - batch_id = tl.load(batch_id_per_token_ptr + token_idx) - local_off = tl.load(local_offset_per_token_ptr + token_idx) - - block_idx = local_off // block_size - block_offset = local_off % block_size - - physical_block_id = tl.load(block_tables_ptr + batch_id * max_blocks_per_seq + block_idx) - - latent_base = latent_cache_ptr + physical_block_id * block_size * LATENT_DIM + block_offset * LATENT_DIM - - kv_c_offs = tl.arange(0, kv_lora_rank) - kv_c_val = tl.load(latent_base + kv_c_offs) - tl.store(output_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) - - k_pe_offs = tl.arange(0, qk_rope_head_dim) - k_pe_val = tl.load(latent_base + kv_lora_rank + k_pe_offs) - tl.store(output_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) - - -def read_latent_from_cache_triton( - latent_cache: paddle.Tensor, - block_tables: paddle.Tensor, - cache_kv_lens: paddle.Tensor, - cu_seqlens_cached_kv: paddle.Tensor, - total_cached_tokens: int, - block_size: int, - kv_lora_rank: int, - qk_rope_head_dim: int, -): - """Triton-accelerated version of read_latent_from_cache_naive. - - Returns (cached_kv_c, cached_k_pe) with the same semantics as the naive - implementation. - """ - if total_cached_tokens == 0: - return None, None - - bsz = cu_seqlens_cached_kv.shape[0] - 1 - max_blocks_per_seq = block_tables.shape[1] - - # --- host-side build of batch_id_per_token / local_offset_per_token --- - # O(bsz) CPU work; much cheaper than the in-kernel linear scan used by the - # legacy read_latent_from_cache_kernel. - cu_list = cu_seqlens_cached_kv.tolist() if hasattr(cu_seqlens_cached_kv, "tolist") else list(cu_seqlens_cached_kv) - batch_id_host = [0] * total_cached_tokens - local_off_host = [0] * total_cached_tokens + out_pos = 0 for b in range(bsz): - start = int(cu_list[b]) - end = int(cu_list[b + 1]) - for t in range(start, end): - batch_id_host[t] = b - local_off_host[t] = t - start - batch_id_per_token = paddle.to_tensor(batch_id_host, dtype="int32", place=latent_cache.place) - local_offset_per_token = paddle.to_tensor(local_off_host, dtype="int32", place=latent_cache.place) - - cached_kv_c = paddle.empty([total_cached_tokens, kv_lora_rank], dtype=latent_cache.dtype) - cached_k_pe = paddle.empty([total_cached_tokens, qk_rope_head_dim], dtype=latent_cache.dtype) - - grid = (total_cached_tokens,) - _read_latent_triton_kernel[grid]( - latent_cache, - block_tables, - batch_id_per_token, - local_offset_per_token, - cached_kv_c, - cached_k_pe, - max_blocks_per_seq=max_blocks_per_seq, - block_size=block_size, - kv_lora_rank=kv_lora_rank, - qk_rope_head_dim=qk_rope_head_dim, - LATENT_DIM=kv_lora_rank + qk_rope_head_dim, - ) + nc = int(cu_cached[b + 1]) - int(cu_cached[b]) + nn = int(cu_new[b + 1]) - int(cu_new[b]) + # cached tokens first + for t in range(nc): + block_idx = t // block_size + block_offset = t % block_size + physical_block_id = block_tables[b, block_idx].item() + latent_vec = latent_cache[physical_block_id, 0, block_offset, :] + full_compressed_kv[out_pos] = latent_vec[:kv_lora_rank] + full_k_pe[out_pos] = latent_vec[kv_lora_rank:] + out_pos += 1 + # new tokens after cached + new_base = int(cu_new[b]) + for t in range(nn): + full_compressed_kv[out_pos] = new_compressed_kv[new_base + t] + full_k_pe[out_pos] = new_k_pe[new_base + t] + out_pos += 1 assert ( - cached_kv_c.shape[0] == total_cached_tokens - ), f"read_latent_from_cache_triton: output shape mismatch {cached_kv_c.shape[0]} vs {total_cached_tokens}" - return cached_kv_c, cached_k_pe - - -def read_latent_from_cache(*args, **kwargs): - if os.environ.get("FD_MLA_USE_NAIVE", "0") == "1": - return read_latent_from_cache_naive(*args, **kwargs) - return read_latent_from_cache_triton(*args, **kwargs) - - -# ---------------------------------------------------------------------------- -# Triton implementation of interleave_cached_and_new_latent. -# ---------------------------------------------------------------------------- + out_pos == total_tokens + ), f"fused_read_cache_and_interleave_naive: out_pos={out_pos} != total_tokens={total_tokens}" + return full_compressed_kv, full_k_pe @enable_compat_on_triton_kernel @triton.jit() -def _interleave_latent_kernel( - cached_kv_c_ptr, # [total_cached, kv_lora_rank] (may be invalid when total_cached==0) - cached_k_pe_ptr, # [total_cached, qk_rope_head_dim] +def _fused_read_interleave_kernel( + latent_cache_ptr, # [num_blocks, 1, block_size, LATENT_DIM] new_kv_c_ptr, # [total_new, kv_lora_rank] new_k_pe_ptr, # [total_new, qk_rope_head_dim] - src_is_cached_ptr, # [total_tokens] int32 (1 == cached, 0 == new) - src_idx_ptr, # [total_tokens] int32 (index within the chosen source tensor) + is_cached_ptr, # [total_tokens] int32 (1 = cached, 0 = new) + src_off_ptr, # [total_tokens] int32 out_kv_c_ptr, # [total_tokens, kv_lora_rank] out_k_pe_ptr, # [total_tokens, qk_rope_head_dim] kv_lora_rank: tl.constexpr, qk_rope_head_dim: tl.constexpr, + LATENT_DIM: tl.constexpr, ): token_idx = tl.program_id(axis=0) - is_cached = tl.load(src_is_cached_ptr + token_idx) - src_idx = tl.load(src_idx_ptr + token_idx) + is_cached = tl.load(is_cached_ptr + token_idx) + src_off = tl.load(src_off_ptr + token_idx) kv_c_offs = tl.arange(0, kv_lora_rank) k_pe_offs = tl.arange(0, qk_rope_head_dim) if is_cached != 0: - kv_c_val = tl.load(cached_kv_c_ptr + src_idx * kv_lora_rank + kv_c_offs) - k_pe_val = tl.load(cached_k_pe_ptr + src_idx * qk_rope_head_dim + k_pe_offs) + # src_off is a flat token offset into latent_cache: physical_block_id * block_size + block_offset + base = latent_cache_ptr + src_off * LATENT_DIM + kv_c_val = tl.load(base + kv_c_offs) + k_pe_val = tl.load(base + kv_lora_rank + k_pe_offs) else: - kv_c_val = tl.load(new_kv_c_ptr + src_idx * kv_lora_rank + kv_c_offs) - k_pe_val = tl.load(new_k_pe_ptr + src_idx * qk_rope_head_dim + k_pe_offs) + kv_c_val = tl.load(new_kv_c_ptr + src_off * kv_lora_rank + kv_c_offs) + k_pe_val = tl.load(new_k_pe_ptr + src_off * qk_rope_head_dim + k_pe_offs) tl.store(out_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) tl.store(out_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) -def interleave_cached_and_new_latent_triton( - cached_kv_c: paddle.Tensor, - cached_k_pe: paddle.Tensor, +def fused_read_cache_and_interleave_triton( + latent_cache: paddle.Tensor, + block_tables: paddle.Tensor, new_compressed_kv: paddle.Tensor, new_k_pe: paddle.Tensor, cu_seqlens_cached_kv: paddle.Tensor, - seq_lens_encoder: paddle.Tensor, cu_seqlens_q: paddle.Tensor, kv_lora_rank: int, qk_rope_head_dim: int, -): - """Triton-accelerated version of interleave_cached_and_new_latent_naive.""" + block_size: int, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + """Triton-accelerated fused read+interleave. + + Host-side builds ``is_cached[t]`` and ``src_off[t]`` (O(total_tokens) CPU + work) so the kernel body contains only a single branch with no table walk. + """ bsz = cu_seqlens_cached_kv.shape[0] - 1 - total_cached = cached_kv_c.shape[0] if cached_kv_c is not None and cached_kv_c.numel() > 0 else 0 + cu_cached = cu_seqlens_cached_kv.tolist() + cu_new = cu_seqlens_q.tolist() + total_cached = int(cu_cached[bsz]) total_new = new_compressed_kv.shape[0] total_tokens = total_cached + total_new full_compressed_kv = paddle.empty([total_tokens, kv_lora_rank], dtype=new_compressed_kv.dtype) full_k_pe = paddle.empty([total_tokens, qk_rope_head_dim], dtype=new_k_pe.dtype) - if total_tokens == 0: return full_compressed_kv, full_k_pe - # ---- host-side mapping: for each output position, pick source ---- - cu_cached = ( - cu_seqlens_cached_kv.tolist() if hasattr(cu_seqlens_cached_kv, "tolist") else list(cu_seqlens_cached_kv) - ) - cu_new = cu_seqlens_q.tolist() if hasattr(cu_seqlens_q, "tolist") else list(cu_seqlens_q) + # block_tables.tolist() is a one-shot D2H; acceptable since host-side loop + # already requires CPU iteration over total_tokens. + bt_list = block_tables.tolist() - src_is_cached = [0] * total_tokens - src_idx = [0] * total_tokens + is_cached_host = [0] * total_tokens + src_off_host = [0] * total_tokens out_pos = 0 for b in range(bsz): nc = int(cu_cached[b + 1]) - int(cu_cached[b]) nn = int(cu_new[b + 1]) - int(cu_new[b]) - cached_base = int(cu_cached[b]) - new_base = int(cu_new[b]) for t in range(nc): - src_is_cached[out_pos] = 1 - src_idx[out_pos] = cached_base + t + block_idx = t // block_size + block_offset = t % block_size + physical_block_id = int(bt_list[b][block_idx]) + is_cached_host[out_pos] = 1 + src_off_host[out_pos] = physical_block_id * block_size + block_offset out_pos += 1 + new_base = int(cu_new[b]) for t in range(nn): - src_is_cached[out_pos] = 0 - src_idx[out_pos] = new_base + t + is_cached_host[out_pos] = 0 + src_off_host[out_pos] = new_base + t out_pos += 1 assert ( out_pos == total_tokens - ), f"interleave_cached_and_new_latent_triton: host map out_pos={out_pos} != total_tokens={total_tokens}" + ), f"fused_read_cache_and_interleave_triton: out_pos={out_pos} != total_tokens={total_tokens}" dev = new_compressed_kv.place - src_is_cached_t = paddle.to_tensor(src_is_cached, dtype="int32", place=dev) - src_idx_t = paddle.to_tensor(src_idx, dtype="int32", place=dev) - - # Provide non-null pointers for the cached tensors even when total_cached == 0; - # Triton kernel path is still gated by src_is_cached flag. - cached_kv_c_ptr = cached_kv_c if total_cached > 0 else full_compressed_kv - cached_k_pe_ptr = cached_k_pe if total_cached > 0 else full_k_pe + is_cached_t = paddle.to_tensor(is_cached_host, dtype="int32", place=dev) + src_off_t = paddle.to_tensor(src_off_host, dtype="int32", place=dev) grid = (total_tokens,) - _interleave_latent_kernel[grid]( - cached_kv_c_ptr, - cached_k_pe_ptr, + _fused_read_interleave_kernel[grid]( + latent_cache, new_compressed_kv, new_k_pe, - src_is_cached_t, - src_idx_t, + is_cached_t, + src_off_t, full_compressed_kv, full_k_pe, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, + LATENT_DIM=kv_lora_rank + qk_rope_head_dim, ) - return full_compressed_kv, full_k_pe -def interleave_cached_and_new_latent(*args, **kwargs): +def fused_read_cache_and_interleave(*args, **kwargs): + """Unified entry. ``FD_MLA_USE_NAIVE=1`` forces the Python reference path.""" if os.environ.get("FD_MLA_USE_NAIVE", "0") == "1": - return interleave_cached_and_new_latent_naive(*args, **kwargs) - return interleave_cached_and_new_latent_triton(*args, **kwargs) - - -# ============================================================================ + return fused_read_cache_and_interleave_naive(*args, **kwargs) + return fused_read_cache_and_interleave_triton(*args, **kwargs) @enable_compat_on_triton_kernel @@ -867,7 +618,8 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): # - new_tokens_of_this_batch = seq_lens_this_time[i] (== enc_len for prefill, == 1 for decode) # - cached_tokens = dec_len if dec_len > 0 else 0 # cu_seqlens_k_with_cache must reflect this sum per batch. - # cu_seqlens_cached_kv tracks only the cached portion for read_latent_from_cache(). + # cu_seqlens_cached_kv tracks only the cached portion consumed by + # fused_read_cache_and_interleave(). for i in range(bsz): # enc_len = ( # forward_meta.seq_lens_encoder[i].item() diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index 78d05aabcb1..9fd88cea884 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -355,8 +355,7 @@ def forward( from fastdeploy.model_executor.layers.attention.mla_attention_backend import ( MLAAttentionMetadata, - interleave_cached_and_new_latent, - read_latent_from_cache, + fused_read_cache_and_interleave, ) attn_out = None @@ -393,50 +392,27 @@ def forward( has_prefix_cache = attn_meta.has_prefix_cache total_cached_tokens = attn_meta.total_cached_kv_tokens - # Handle prefix cache: read and concatenate cached latent with new latent + # Handle prefix cache: read cached latent from paged cache and interleave + # with the new-token latent in a single fused kernel call. + full_compressed_kv = compressed_kv + full_k_pe = key_pe.squeeze(1) if has_prefix_cache and total_cached_tokens > 0: - # Get latent cache from forward_meta layer_id = self.mla_attn.layer_id if hasattr(self.mla_attn, "layer_id") else 0 latent_cache = forward_meta.caches[layer_id] if hasattr(forward_meta, "caches") else None - if latent_cache is not None: - # Read cached latent vectors - cached_kv_c, cached_k_pe = read_latent_from_cache( + block_size = self.mla_attn.block_size if hasattr(self.mla_attn, "block_size") else 64 + full_compressed_kv, full_k_pe = fused_read_cache_and_interleave( latent_cache, forward_meta.block_tables, - forward_meta.seq_lens_decoder, # cache_kv_lens + compressed_kv, + key_pe.squeeze(1), attn_meta.cu_seqlens_cached_kv, - total_cached_tokens, - self.mla_attn.block_size if hasattr(self.mla_attn, "block_size") else 64, + forward_meta.cu_seqlens_q, self.kv_lora_rank, self.qk_rope_head_dim, + block_size, ) - if cached_kv_c is not None and cached_k_pe is not None: - # Interleave cached and new latent vectors per batch - # This ensures correct ordering: [batch0_cached, batch0_new, batch1_cached, batch1_new, ...] - # cached_k_pe already has RoPE applied (stored with RoPE) - full_compressed_kv, full_k_pe = interleave_cached_and_new_latent( - cached_kv_c, - cached_k_pe, - compressed_kv, - key_pe.squeeze(1), - attn_meta.cu_seqlens_cached_kv, - forward_meta.seq_lens_encoder, - forward_meta.cu_seqlens_q, - self.kv_lora_rank, - self.qk_rope_head_dim, - ) - else: - full_compressed_kv = compressed_kv - full_k_pe = key_pe.squeeze(1) - else: - full_compressed_kv = compressed_kv - full_k_pe = key_pe.squeeze(1) - else: - full_compressed_kv = compressed_kv - full_k_pe = key_pe.squeeze(1) - # Project latent KV to full key and value key_value = self.kv_b_proj(full_compressed_kv) key_value.reshape_( From b4083fa778a6ea0a031336191583303489637f8a Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Fri, 8 May 2026 15:49:50 +0800 Subject: [PATCH 05/10] support mla chunk-prefill & prefix cache --- .../layers/attention/mla_attention_backend.py | 153 ++++++------------ 1 file changed, 52 insertions(+), 101 deletions(-) diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 050e3e90119..daf736bd86a 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -559,113 +559,64 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): metadata.max_dec_len_this_time = forward_meta.max_len_tensor_cpu[2] metadata.max_kv_len_this_time = forward_meta.max_len_tensor_cpu[5] - # Compute prefix cache metadata - # For MLA, prefix cache is indicated by seq_lens_decoder > 0 when there are encoder tokens - # This means the request has cached KV from a previous prefix - bsz = forward_meta.seq_lens_this_time.shape[0] - has_prefix_cache = False - total_cached_kv_tokens = 0 - max_cached_kv_len = 0 - max_total_kv_len = 0 # max(dec_len + enc_len) for FlashAttention - - # Check if any request has prefix cache - # Prefix cache exists when seq_lens_decoder > 0 - # seq_lens_decoder stores the cached KV length for chunked prefill/prefix cache - for i in range(bsz): - # enc_len = ( - # forward_meta.seq_lens_encoder[i].item() - # if hasattr(forward_meta.seq_lens_encoder[i], "item") - # else forward_meta.seq_lens_encoder[i] - # ) - dec_len = ( - forward_meta.seq_lens_decoder[i].item() - if hasattr(forward_meta.seq_lens_decoder[i], "item") - else forward_meta.seq_lens_decoder[i] - ) - seq_this_time = ( - forward_meta.seq_lens_this_time[i].item() - if hasattr(forward_meta.seq_lens_this_time[i], "item") - else forward_meta.seq_lens_this_time[i] - ) - # Per-batch K length after interleave = dec_len (cached) + seq_this_time (new). - # seq_this_time equals enc_len for prefill/chunked-prefill, and 1 for decode. - # Using enc_len here drops the decode-batch new token and yields a wrong - # max_seqlen_k for FlashAttention, which then tile-truncates the last K row. - per_batch_k = dec_len + seq_this_time - if dec_len > 0: - has_prefix_cache = True - total_cached_kv_tokens += dec_len - max_cached_kv_len = max(max_cached_kv_len, dec_len) - if per_batch_k > 0: - max_total_kv_len = max(max_total_kv_len, per_batch_k) - - metadata.has_prefix_cache = has_prefix_cache - metadata.total_cached_kv_tokens = total_cached_kv_tokens - metadata.max_cached_kv_len = max_cached_kv_len - metadata.max_total_kv_len = max_total_kv_len - - # Compute cu_seqlens_cached_kv if there's prefix cache - if has_prefix_cache and total_cached_kv_tokens > 0: - cu_seqlens_cached_kv = paddle.zeros([bsz + 1], dtype=paddle.int32) - cu_seqlens_k_with_cache = paddle.zeros([bsz + 1], dtype=paddle.int32) + # Prefix-cache metadata build. + # Fast path: max(seq_lens_decoder)==0 ⇔ no prefix cache this step. Reading from + # max_len_tensor_cpu (already CPU, populated by get_block_shape_and_split_kv_block) + # costs zero D2H, skipping all loops/H2D below. + if int(forward_meta.max_len_tensor_cpu[2]) == 0: + metadata.has_prefix_cache = False + metadata.total_cached_kv_tokens = 0 + metadata.max_cached_kv_len = 0 + metadata.max_total_kv_len = 0 # unused when has_prefix_cache=False + else: + # Slow path: per-batch K layout after interleave = [dec_len cached, seq_this_time new]. + bsz = forward_meta.seq_lens_this_time.shape[0] + # Single batched D2H for both length tensors. + stacked = paddle.stack([forward_meta.seq_lens_decoder, forward_meta.seq_lens_this_time]).tolist() + dec_lens, seq_this_times = stacked[0], stacked[1] + + total_cached_kv_tokens = 0 + max_cached_kv_len = 0 + max_total_kv_len = 0 + cu_cached = [0] * (bsz + 1) + cu_total = [0] * (bsz + 1) cumsum_cached = 0 cumsum_total = 0 - # cu_seqlens layout must stay CONSISTENT with the input tensor layout used - # by the prefill FlashAttention call. The input `compressed_kv`/`key_pe` contains - # ALL tokens in the batch (prefill + decode), in seq_lens_this_time order. - # - # For each batch i, the interleaved layout writes: [cached_tokens, new_tokens_of_this_batch] - # - new_tokens_of_this_batch = seq_lens_this_time[i] (== enc_len for prefill, == 1 for decode) - # - cached_tokens = dec_len if dec_len > 0 else 0 - # cu_seqlens_k_with_cache must reflect this sum per batch. - # cu_seqlens_cached_kv tracks only the cached portion consumed by - # fused_read_cache_and_interleave(). for i in range(bsz): - # enc_len = ( - # forward_meta.seq_lens_encoder[i].item() - # if hasattr(forward_meta.seq_lens_encoder[i], "item") - # else forward_meta.seq_lens_encoder[i] - # ) - dec_len = ( - forward_meta.seq_lens_decoder[i].item() - if hasattr(forward_meta.seq_lens_decoder[i], "item") - else forward_meta.seq_lens_decoder[i] - ) - seq_this_time = ( - forward_meta.seq_lens_this_time[i].item() - if hasattr(forward_meta.seq_lens_this_time[i], "item") - else forward_meta.seq_lens_this_time[i] - ) + dec_len = int(dec_lens[i]) + seq_this = int(seq_this_times[i]) if dec_len > 0: + total_cached_kv_tokens += dec_len + if dec_len > max_cached_kv_len: + max_cached_kv_len = dec_len cumsum_cached += dec_len cumsum_total += dec_len - # Add this batch's new tokens to cumsum_total. Use seq_lens_this_time to cover - # both prefill (== enc_len) and decode (== 1) correctly. - cumsum_total += seq_this_time - cu_seqlens_cached_kv[i + 1] = cumsum_cached - cu_seqlens_k_with_cache[i + 1] = cumsum_total - # Consistency checks: starts at 0, monotonic non-decreasing, final equals cumulative. - assert cu_seqlens_cached_kv[0].item() == 0, "cu_seqlens_cached_kv must start at 0" - assert cu_seqlens_k_with_cache[0].item() == 0, "cu_seqlens_k_with_cache must start at 0" - assert ( - cu_seqlens_cached_kv[bsz].item() == cumsum_cached - ), f"cu_seqlens_cached_kv[-1]={cu_seqlens_cached_kv[bsz].item()} != cumsum_cached={cumsum_cached}" - assert ( - cu_seqlens_k_with_cache[bsz].item() == cumsum_total - ), f"cu_seqlens_k_with_cache[-1]={cu_seqlens_k_with_cache[bsz].item()} != cumsum_total={cumsum_total}" - metadata.cu_seqlens_cached_kv = cu_seqlens_cached_kv - metadata.cu_seqlens_k_with_cache = cu_seqlens_k_with_cache - # Cross-check: max_total_kv_len MUST be >= max per-batch K length implied by - # cu_seqlens_k_with_cache. Otherwise FlashAttention will tile-truncate the - # last K rows of the longest batch and silently corrupt attention output. - if bsz > 0: - cu_k_list = cu_seqlens_k_with_cache.tolist() - observed_max = max(cu_k_list[i + 1] - cu_k_list[i] for i in range(bsz)) - assert max_total_kv_len >= observed_max, ( - f"max_total_kv_len={max_total_kv_len} < observed per-batch max " - f"K length={observed_max} from cu_seqlens_k_with_cache={cu_k_list}. " - f"FlashAttention will truncate K tiles." - ) + cumsum_total += seq_this + per_batch_k = dec_len + seq_this + if per_batch_k > max_total_kv_len: + max_total_kv_len = per_batch_k + cu_cached[i + 1] = cumsum_cached + cu_total[i + 1] = cumsum_total + + # Entering slow path implies max_dec>0 ⇒ at least one batch has dec_len>0. + metadata.has_prefix_cache = True + metadata.total_cached_kv_tokens = total_cached_kv_tokens + metadata.max_cached_kv_len = max_cached_kv_len + metadata.max_total_kv_len = max_total_kv_len + + # Single batched H2D: [2, bsz+1] → two zero-copy views. + dev = forward_meta.seq_lens_decoder.place + cu_pair = paddle.to_tensor([cu_cached, cu_total], dtype=paddle.int32, place=dev) + metadata.cu_seqlens_cached_kv = cu_pair[0] + metadata.cu_seqlens_k_with_cache = cu_pair[1] + + # Cross-check precision guard: max_total_kv_len MUST dominate per-batch K length, + # otherwise FlashAttention tile-truncates the last K rows and silently corrupts output. + observed_max = max(cu_total[i + 1] - cu_total[i] for i in range(bsz)) if bsz > 0 else 0 + assert max_total_kv_len >= observed_max, ( + f"max_total_kv_len={max_total_kv_len} < observed per-batch max " + f"K length={observed_max}. FlashAttention would truncate K tiles." + ) # pd_disaggregation metadata.kv_signal_data_list = [None] * self.num_layers From 9c8c022e799be881e2c62847e6a16ecc10b72bb5 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Fri, 8 May 2026 22:05:55 +0800 Subject: [PATCH 06/10] update mla chunk prefill prefix cache --- custom_ops/gpu_ops/cpp_extensions.cc | 1 + custom_ops/gpu_ops/get_padding_offset.cu | 39 ++- .../layers/attention/mla_attention_backend.py | 263 ++++++------------ .../model_executor/models/deepseek_v3.py | 42 ++- .../model_executor/pre_and_post_process.py | 2 +- 5 files changed, 128 insertions(+), 219 deletions(-) diff --git a/custom_ops/gpu_ops/cpp_extensions.cc b/custom_ops/gpu_ops/cpp_extensions.cc index 204ea33e50b..d7f53c46209 100644 --- a/custom_ops/gpu_ops/cpp_extensions.cc +++ b/custom_ops/gpu_ops/cpp_extensions.cc @@ -427,6 +427,7 @@ std::vector GetPaddingOffset( const paddle::Tensor& seq_len, const paddle::optional& draft_tokens, const paddle::optional& seq_lens_encoder, + const paddle::optional& seq_lens_decoder, const int64_t token_num_cpu); void SetValueByFlagsAndIdx(const paddle::Tensor& token_ids_all, diff --git a/custom_ops/gpu_ops/get_padding_offset.cu b/custom_ops/gpu_ops/get_padding_offset.cu index 998d1760515..010eb709fe5 100644 --- a/custom_ops/gpu_ops/get_padding_offset.cu +++ b/custom_ops/gpu_ops/get_padding_offset.cu @@ -28,6 +28,7 @@ __global__ void PrefixSumKernel(int64_t *ids_remove_padding, const int max_seq_len, const int64_t *draft_tokens, const int *seq_lens_encoder, + const int *seq_lens_decoder, const int max_draft_tokens_per_batch) { const int bi = blockIdx.x; const int tid = threadIdx.x; @@ -39,23 +40,37 @@ __global__ void PrefixSumKernel(int64_t *ids_remove_padding, const int lane_id = threadIdx.x % 32; #endif - int cum_seq_len = 0; + // Independent q/k prefix sums: + // cu_seqlens_q = Σ seq_lens[j] (new tokens only, for Q/varlen indexing) + // cu_seqlens_k = Σ (seq_lens[j] + seq_lens_decoder[j]) (cached + new, for + // K/FA) + // When seq_lens_decoder == nullptr, cu_seqlens_k degenerates to cu_seqlens_q. + int cum_seq_len_q = 0; + int cum_seq_len_k = 0; - // compute sum of seq_lens[0,1,2,...,bi] for (int i = lane_id; i < bi + 1; i += WARP_SIZE) { - cum_seq_len += seq_lens[i]; + const int q_inc = seq_lens[i]; + const int k_inc = + q_inc + (seq_lens_decoder != nullptr ? seq_lens_decoder[i] : 0); + cum_seq_len_q += q_inc; + cum_seq_len_k += k_inc; } for (int offset = 1; offset < WARP_SIZE; offset <<= 1) { - const int tmp = __shfl_up_sync(0xffffffff, cum_seq_len, offset); - if (lane_id >= offset) cum_seq_len += tmp; + const int tmp_q = __shfl_up_sync(0xffffffff, cum_seq_len_q, offset); + const int tmp_k = __shfl_up_sync(0xffffffff, cum_seq_len_k, offset); + if (lane_id >= offset) { + cum_seq_len_q += tmp_q; + cum_seq_len_k += tmp_k; + } } - cum_seq_len = __shfl_sync(0xffffffff, cum_seq_len, WARP_SIZE - 1); + cum_seq_len_q = __shfl_sync(0xffffffff, cum_seq_len_q, WARP_SIZE - 1); + cum_seq_len_k = __shfl_sync(0xffffffff, cum_seq_len_k, WARP_SIZE - 1); if (tid == 0) { - cu_seqlens_q[bi + 1] = cum_seq_len; - cu_seqlens_k[bi + 1] = cum_seq_len; + cu_seqlens_q[bi + 1] = cum_seq_len_q; + cu_seqlens_k[bi + 1] = cum_seq_len_k; } if (bi == 0 && tid == 0) { @@ -63,8 +78,9 @@ __global__ void PrefixSumKernel(int64_t *ids_remove_padding, cu_seqlens_k[0] = 0; } + // Q-side token scatter uses cum_seq_len_q (new-tokens-only layout). for (int i = tid; i < seq_lens[bi]; i += blockDim.x) { - const int tgt_seq_id = cum_seq_len - seq_lens[bi] + i; + const int tgt_seq_id = cum_seq_len_q - seq_lens[bi] + i; if (max_draft_tokens_per_batch > 0 && seq_lens_encoder[bi] <= 0) { // speculative decoding const int src_seq_id = bi * max_draft_tokens_per_batch + i; @@ -84,6 +100,7 @@ std::vector GetPaddingOffset( const paddle::Tensor &seq_len, const paddle::optional &draft_tokens, const paddle::optional &seq_lens_encoder, + const paddle::optional &seq_lens_decoder, const int64_t cpu_token_num) { #ifdef PADDLE_WITH_CUSTOM_DEVICE auto dev_ctx = static_cast( @@ -131,6 +148,7 @@ std::vector GetPaddingOffset( max_seq_len, draft_tokens ? draft_tokens.get().data() : nullptr, seq_lens_encoder ? seq_lens_encoder.get().data() : nullptr, + seq_lens_decoder ? seq_lens_decoder.get().data() : nullptr, max_draft_tokens_per_batch); return {x_remove_padding, batch_id_per_token, cu_seqlens_q, cu_seqlens_k}; @@ -156,7 +174,8 @@ PD_BUILD_STATIC_OP(get_padding_offset) .Inputs({"input_ids", "seq_len", paddle::Optional("draft_tokens"), - paddle::Optional("seq_lens_encoder")}) + paddle::Optional("seq_lens_encoder"), + paddle::Optional("seq_lens_decoder")}) .Outputs({"x_remove_padding", "batch_id_per_token", "cu_seqlens_q", diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index daf736bd86a..ede903f32f0 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -83,7 +83,7 @@ def fused_read_cache_and_interleave_naive( block_tables: paddle.Tensor, new_compressed_kv: paddle.Tensor, new_k_pe: paddle.Tensor, - cu_seqlens_cached_kv: paddle.Tensor, + cu_seqlens_k: paddle.Tensor, cu_seqlens_q: paddle.Tensor, kv_lora_rank: int, qk_rope_head_dim: int, @@ -91,15 +91,14 @@ def fused_read_cache_and_interleave_naive( ) -> Tuple[paddle.Tensor, paddle.Tensor]: """Python/Paddle reference of the fused read+interleave path. - Returns ``(full_compressed_kv, full_k_pe)`` with per-batch layout - ``[cached_tokens, new_tokens]`` in token order. + Takes with-cache ``cu_seqlens_k`` (cumsum of ``cached + new`` + per batch) plus ``cu_seqlens_q`` (new-only); the cached length of each + batch is derived as ``k_len - new_len``. """ - bsz = cu_seqlens_cached_kv.shape[0] - 1 - cu_cached = cu_seqlens_cached_kv.tolist() + bsz = cu_seqlens_k.shape[0] - 1 + cu_total = cu_seqlens_k.tolist() cu_new = cu_seqlens_q.tolist() - total_cached = int(cu_cached[bsz]) - total_new = new_compressed_kv.shape[0] - total_tokens = total_cached + total_new + total_tokens = int(cu_total[bsz]) full_compressed_kv = paddle.empty([total_tokens, kv_lora_rank], dtype=new_compressed_kv.dtype) full_k_pe = paddle.empty([total_tokens, qk_rope_head_dim], dtype=new_k_pe.dtype) @@ -108,8 +107,9 @@ def fused_read_cache_and_interleave_naive( out_pos = 0 for b in range(bsz): - nc = int(cu_cached[b + 1]) - int(cu_cached[b]) + k_len = int(cu_total[b + 1]) - int(cu_total[b]) nn = int(cu_new[b + 1]) - int(cu_new[b]) + nc = k_len - nn # cached tokens first for t in range(nc): block_idx = t // block_size @@ -138,29 +138,62 @@ def _fused_read_interleave_kernel( latent_cache_ptr, # [num_blocks, 1, block_size, LATENT_DIM] new_kv_c_ptr, # [total_new, kv_lora_rank] new_k_pe_ptr, # [total_new, qk_rope_head_dim] - is_cached_ptr, # [total_tokens] int32 (1 = cached, 0 = new) - src_off_ptr, # [total_tokens] int32 + cu_total_ptr, # [bsz+1] int32 (= forward_meta.cu_seqlens_k) + cu_new_ptr, # [bsz+1] int32 (= cu_seqlens_q) + block_tables_ptr, # [bsz, max_blocks_per_seq] int32 out_kv_c_ptr, # [total_tokens, kv_lora_rank] out_k_pe_ptr, # [total_tokens, qk_rope_head_dim] + bsz, + max_blocks_per_seq: tl.constexpr, + block_size: tl.constexpr, kv_lora_rank: tl.constexpr, qk_rope_head_dim: tl.constexpr, LATENT_DIM: tl.constexpr, ): + """Single-kernel fused read+interleave. + + Output tokens are laid out contiguously in batch-major order with per-batch + layout ``[cached, new]``. The owning batch ``b`` for each ``token_idx`` is + located via a linear scan on ``cu_total`` (bsz is small; O(bsz) overhead is + negligible next to per-token LATENT_DIM load/store). Per-batch cached + length is derived in-kernel as + ``nc = (cu_total[b+1] - cu_total[b]) - (cu_new[b+1] - cu_new[b])``. + """ token_idx = tl.program_id(axis=0) - is_cached = tl.load(is_cached_ptr + token_idx) - src_off = tl.load(src_off_ptr + token_idx) + + # Linear scan: largest b such that cu_total[b] <= token_idx. + b = 0 + for i in range(bsz): + base_i = tl.load(cu_total_ptr + i) + if token_idx >= base_i: + b = i + + ct_b = tl.load(cu_total_ptr + b) + ct_b1 = tl.load(cu_total_ptr + b + 1) + cn_b = tl.load(cu_new_ptr + b) + cn_b1 = tl.load(cu_new_ptr + b + 1) + + k_len = ct_b1 - ct_b + nn = cn_b1 - cn_b + nc = k_len - nn + local_t = token_idx - ct_b kv_c_offs = tl.arange(0, kv_lora_rank) k_pe_offs = tl.arange(0, qk_rope_head_dim) - if is_cached != 0: - # src_off is a flat token offset into latent_cache: physical_block_id * block_size + block_offset - base = latent_cache_ptr + src_off * LATENT_DIM + if local_t < nc: + # cached token: walk paged latent_cache via block_tables + block_idx = local_t // block_size + block_off = local_t % block_size + pbid = tl.load(block_tables_ptr + b * max_blocks_per_seq + block_idx) + base = latent_cache_ptr + (pbid * block_size + block_off) * LATENT_DIM kv_c_val = tl.load(base + kv_c_offs) k_pe_val = tl.load(base + kv_lora_rank + k_pe_offs) else: - kv_c_val = tl.load(new_kv_c_ptr + src_off * kv_lora_rank + kv_c_offs) - k_pe_val = tl.load(new_k_pe_ptr + src_off * qk_rope_head_dim + k_pe_offs) + # new token: linear region in new_compressed_kv / new_k_pe + src = cn_b + (local_t - nc) + kv_c_val = tl.load(new_kv_c_ptr + src * kv_lora_rank + kv_c_offs) + k_pe_val = tl.load(new_k_pe_ptr + src * qk_rope_head_dim + k_pe_offs) tl.store(out_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) tl.store(out_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) @@ -171,7 +204,7 @@ def fused_read_cache_and_interleave_triton( block_tables: paddle.Tensor, new_compressed_kv: paddle.Tensor, new_k_pe: paddle.Tensor, - cu_seqlens_cached_kv: paddle.Tensor, + cu_seqlens_k: paddle.Tensor, cu_seqlens_q: paddle.Tensor, kv_lora_rank: int, qk_rope_head_dim: int, @@ -179,61 +212,33 @@ def fused_read_cache_and_interleave_triton( ) -> Tuple[paddle.Tensor, paddle.Tensor]: """Triton-accelerated fused read+interleave. - Host-side builds ``is_cached[t]`` and ``src_off[t]`` (O(total_tokens) CPU - work) so the kernel body contains only a single branch with no table walk. + Single kernel launch with zero Python metadata computation. All per-batch + geometry is derived in-kernel from ``cu_seqlens_k`` and + ``cu_seqlens_q``. Only one scalar D2H is needed: the final cumsum entry, + used to size the output tensors. """ - bsz = cu_seqlens_cached_kv.shape[0] - 1 - cu_cached = cu_seqlens_cached_kv.tolist() - cu_new = cu_seqlens_q.tolist() - total_cached = int(cu_cached[bsz]) - total_new = new_compressed_kv.shape[0] - total_tokens = total_cached + total_new + bsz = cu_seqlens_k.shape[0] - 1 + total_tokens = int(cu_seqlens_k[-1]) full_compressed_kv = paddle.empty([total_tokens, kv_lora_rank], dtype=new_compressed_kv.dtype) full_k_pe = paddle.empty([total_tokens, qk_rope_head_dim], dtype=new_k_pe.dtype) if total_tokens == 0: return full_compressed_kv, full_k_pe - # block_tables.tolist() is a one-shot D2H; acceptable since host-side loop - # already requires CPU iteration over total_tokens. - bt_list = block_tables.tolist() + max_blocks_per_seq = block_tables.shape[1] - is_cached_host = [0] * total_tokens - src_off_host = [0] * total_tokens - out_pos = 0 - for b in range(bsz): - nc = int(cu_cached[b + 1]) - int(cu_cached[b]) - nn = int(cu_new[b + 1]) - int(cu_new[b]) - for t in range(nc): - block_idx = t // block_size - block_offset = t % block_size - physical_block_id = int(bt_list[b][block_idx]) - is_cached_host[out_pos] = 1 - src_off_host[out_pos] = physical_block_id * block_size + block_offset - out_pos += 1 - new_base = int(cu_new[b]) - for t in range(nn): - is_cached_host[out_pos] = 0 - src_off_host[out_pos] = new_base + t - out_pos += 1 - - assert ( - out_pos == total_tokens - ), f"fused_read_cache_and_interleave_triton: out_pos={out_pos} != total_tokens={total_tokens}" - - dev = new_compressed_kv.place - is_cached_t = paddle.to_tensor(is_cached_host, dtype="int32", place=dev) - src_off_t = paddle.to_tensor(src_off_host, dtype="int32", place=dev) - - grid = (total_tokens,) - _fused_read_interleave_kernel[grid]( + _fused_read_interleave_kernel[(total_tokens,)]( latent_cache, new_compressed_kv, new_k_pe, - is_cached_t, - src_off_t, + cu_seqlens_k, + cu_seqlens_q, + block_tables, full_compressed_kv, full_k_pe, + bsz, + max_blocks_per_seq=max_blocks_per_seq, + block_size=block_size, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, LATENT_DIM=kv_lora_rank + qk_rope_head_dim, @@ -413,18 +418,9 @@ class MLAAttentionMetadata(AttentionMetadata): max_kv_len_this_time: Optional[paddle.Tensor] = None # For prefix cache and chunked prefill support - # Indicates whether any request has prefix cache (cached KV from previous requests) - has_prefix_cache: bool = False - # Total number of cached KV tokens across all requests with prefix cache - total_cached_kv_tokens: int = 0 - # cu_seqlens for cached KV tokens (similar to cu_seqlens_k but for cached portion) - cu_seqlens_cached_kv: Optional[paddle.Tensor] = None - # Maximum cached KV length across all requests - max_cached_kv_len: int = 0 - # cu_seqlens_k that includes cached tokens (for FlashAttention when prefix cache is present) - cu_seqlens_k_with_cache: Optional[paddle.Tensor] = None - # Maximum total KV length (cached + new) across all requests for FlashAttention - max_total_kv_len: int = 0 + # ``cu_seqlens_k`` semantics produced by ``get_padding_offset``. + # forward_meta.cu_seqlens_k: Optional[paddle.Tensor] = None + max_seqlen_k: int = 0 class MLAAttentionBackend(AttentionBackend): @@ -558,65 +554,7 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): metadata.max_enc_len_this_time = forward_meta.max_len_tensor_cpu[1] metadata.max_dec_len_this_time = forward_meta.max_len_tensor_cpu[2] metadata.max_kv_len_this_time = forward_meta.max_len_tensor_cpu[5] - - # Prefix-cache metadata build. - # Fast path: max(seq_lens_decoder)==0 ⇔ no prefix cache this step. Reading from - # max_len_tensor_cpu (already CPU, populated by get_block_shape_and_split_kv_block) - # costs zero D2H, skipping all loops/H2D below. - if int(forward_meta.max_len_tensor_cpu[2]) == 0: - metadata.has_prefix_cache = False - metadata.total_cached_kv_tokens = 0 - metadata.max_cached_kv_len = 0 - metadata.max_total_kv_len = 0 # unused when has_prefix_cache=False - else: - # Slow path: per-batch K layout after interleave = [dec_len cached, seq_this_time new]. - bsz = forward_meta.seq_lens_this_time.shape[0] - # Single batched D2H for both length tensors. - stacked = paddle.stack([forward_meta.seq_lens_decoder, forward_meta.seq_lens_this_time]).tolist() - dec_lens, seq_this_times = stacked[0], stacked[1] - - total_cached_kv_tokens = 0 - max_cached_kv_len = 0 - max_total_kv_len = 0 - cu_cached = [0] * (bsz + 1) - cu_total = [0] * (bsz + 1) - cumsum_cached = 0 - cumsum_total = 0 - for i in range(bsz): - dec_len = int(dec_lens[i]) - seq_this = int(seq_this_times[i]) - if dec_len > 0: - total_cached_kv_tokens += dec_len - if dec_len > max_cached_kv_len: - max_cached_kv_len = dec_len - cumsum_cached += dec_len - cumsum_total += dec_len - cumsum_total += seq_this - per_batch_k = dec_len + seq_this - if per_batch_k > max_total_kv_len: - max_total_kv_len = per_batch_k - cu_cached[i + 1] = cumsum_cached - cu_total[i + 1] = cumsum_total - - # Entering slow path implies max_dec>0 ⇒ at least one batch has dec_len>0. - metadata.has_prefix_cache = True - metadata.total_cached_kv_tokens = total_cached_kv_tokens - metadata.max_cached_kv_len = max_cached_kv_len - metadata.max_total_kv_len = max_total_kv_len - - # Single batched H2D: [2, bsz+1] → two zero-copy views. - dev = forward_meta.seq_lens_decoder.place - cu_pair = paddle.to_tensor([cu_cached, cu_total], dtype=paddle.int32, place=dev) - metadata.cu_seqlens_cached_kv = cu_pair[0] - metadata.cu_seqlens_k_with_cache = cu_pair[1] - - # Cross-check precision guard: max_total_kv_len MUST dominate per-batch K length, - # otherwise FlashAttention tile-truncates the last K rows and silently corrupts output. - observed_max = max(cu_total[i + 1] - cu_total[i] for i in range(bsz)) if bsz > 0 else 0 - assert max_total_kv_len >= observed_max, ( - f"max_total_kv_len={max_total_kv_len} < observed per-batch max " - f"K length={observed_max}. FlashAttention would truncate K tiles." - ) + metadata.max_seqlen_k = max(metadata.max_kv_len_this_time.item(), metadata.max_enc_len_this_time.item()) # pd_disaggregation metadata.kv_signal_data_list = [None] * self.num_layers @@ -664,10 +602,10 @@ def forward_extend( forward_meta: ForwardMeta, ) -> paddle.Tensor: """ - Prefill阶段的前向传播，支持 prefix cache + Prefill阶段的前向传播，支持 chunk prefill /prefix cache - 对于 MLA 模型的 prefix cache 支持： - 1. 如果存在 prefix cache (metadata.has_prefix_cache = True) + 对于 MLA 模型的 chunk prefill /prefix cache 支持： + 1. 如果开启 chunk prefill /prefix cache - k 和 v 应该已经包含了 cached KV 和 new KV 的拼接 - cu_seqlens_k 应该已经调整为包含 cached tokens 2. 如果不存在 prefix cache，行为与之前相同 @@ -697,29 +635,14 @@ def forward_extend( getattr(forward_meta, "max_input_length", -1), ) - # Flash注意力计算 - # 对于 prefix cache 场景： - # - k 和 v 应该已经包含了 cached + new tokens - # - cu_seqlens_k 应该已经调整 - # - max_seqlen_k 应该是 cached_len + new_len 的最大值 - - # 获取正确的 cu_seqlens_k 和 max_seqlen_k - if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: - # When prefix cache is present, use cu_seqlens_k that includes cached tokens - cu_seqlens_k = metadata.cu_seqlens_k_with_cache - max_seqlen_k = metadata.max_total_kv_len # max(dec_len + enc_len) - else: - cu_seqlens_k = forward_meta.cu_seqlens_k - max_seqlen_k = metadata.max_enc_len_this_time - fmha_out = self.flash_attn_func( q, k, v, forward_meta.cu_seqlens_q, - cu_seqlens_k, - metadata.max_enc_len_this_time, # max_seqlen_q - only new tokens - max_seqlen_k, # max_seqlen_k - may include cached tokens + forward_meta.cu_seqlens_k, + metadata.max_enc_len_this_time, + metadata.max_seqlen_k, causal=self.causal, **self.flash_attn_kwargs, )[0] @@ -826,9 +749,9 @@ def forward_mixed( forward_meta: ForwardMeta, ) -> paddle.Tensor: """ - Mixed模式的前向传播，支持 prefix cache + Mixed模式的前向传播，支持 chunk prefill /prefix cache - 对于 MLA 模型的 prefix cache 支持： + 对于 MLA 模型的 chunk prefill /prefix cache 支持： 1. Prefill 分支：k 和 v 应该已包含 cached + new tokens 2. Decode 分支：保持原有 latent attention 逻辑 """ @@ -846,13 +769,6 @@ def forward_mixed( # Prefill branch: k is not None if k is not None: - bsz = forward_meta.cu_seqlens_q.shape[0] - 1 - - # Write cache only for new tokens of prefill/chunked-prefill batches. - # Decode batches (seq_lens_encoder == 0) are intentionally skipped here — they - # are handled later by decode_mla_write_cache() via the k=None branch when - # need_do_decode is True. Using seq_lens_encoder keeps that separation correct - # and avoids double-writing decode tokens. prefill_mla_write_cache( compressed_kv, k_pe, @@ -867,32 +783,15 @@ def forward_mixed( self.max_seq_len, ) - # Determine cu_seqlens_k and max_seqlen_k considering prefix cache - if metadata.has_prefix_cache and metadata.cu_seqlens_k_with_cache is not None: - cu_seqlens_k = metadata.cu_seqlens_k_with_cache - max_seqlen_k = metadata.max_total_kv_len # max(dec_len + enc_len) - else: - cu_seqlens_k = forward_meta.cu_seqlens_k - max_seqlen_k = metadata.max_enc_len_this_time - - # Shape consistency: k/v token count must match cu_seqlens_k terminal value - _expected_k = cu_seqlens_k[bsz].item() if hasattr(cu_seqlens_k[bsz], "item") else int(cu_seqlens_k[bsz]) - assert ( - k.shape[0] == _expected_k - ), f"forward_mixed: k.shape[0]={k.shape[0]} != cu_seqlens_k[-1]={_expected_k}" - assert ( - v.shape[0] == _expected_k - ), f"forward_mixed: v.shape[0]={v.shape[0]} != cu_seqlens_k[-1]={_expected_k}" - # FlashAttention for prefill fmha_out = self.flash_attn_func( q, k, v, forward_meta.cu_seqlens_q, - cu_seqlens_k, - metadata.max_enc_len_this_time, # max_seqlen_q - only new tokens - max_seqlen_k, # max_seqlen_k - may include cached tokens + forward_meta.cu_seqlens_k, + metadata.max_enc_len_this_time, + metadata.max_seqlen_k, causal=self.causal, **self.flash_attn_kwargs, )[0] diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index 9fd88cea884..d75af311546 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -208,6 +208,7 @@ def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None super().__init__() self.fd_config = fd_config + self.layer_id = layer_id self.use_gated_attn = getattr(self.fd_config.model_config, "use_gated_attn", False) self.use_bias = getattr(self.fd_config.model_config, "use_bias", False) self.tp_size = fd_config.parallel_config.tensor_parallel_size @@ -354,7 +355,6 @@ def forward( """MLA attention forward with prefix cache support.""" from fastdeploy.model_executor.layers.attention.mla_attention_backend import ( - MLAAttentionMetadata, fused_read_cache_and_interleave, ) @@ -383,35 +383,25 @@ def forward( need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0 if need_do_prefill: # max_enc_len_this_time - # Check for prefix cache - attn_meta = forward_meta.attn_backend.attention_metadata if hasattr(forward_meta, "attn_backend") else None - has_prefix_cache = False - total_cached_tokens = 0 - - if attn_meta is not None and isinstance(attn_meta, MLAAttentionMetadata): - has_prefix_cache = attn_meta.has_prefix_cache - total_cached_tokens = attn_meta.total_cached_kv_tokens - # Handle prefix cache: read cached latent from paged cache and interleave # with the new-token latent in a single fused kernel call. full_compressed_kv = compressed_kv full_k_pe = key_pe.squeeze(1) - if has_prefix_cache and total_cached_tokens > 0: - layer_id = self.mla_attn.layer_id if hasattr(self.mla_attn, "layer_id") else 0 - latent_cache = forward_meta.caches[layer_id] if hasattr(forward_meta, "caches") else None - if latent_cache is not None: - block_size = self.mla_attn.block_size if hasattr(self.mla_attn, "block_size") else 64 - full_compressed_kv, full_k_pe = fused_read_cache_and_interleave( - latent_cache, - forward_meta.block_tables, - compressed_kv, - key_pe.squeeze(1), - attn_meta.cu_seqlens_cached_kv, - forward_meta.cu_seqlens_q, - self.kv_lora_rank, - self.qk_rope_head_dim, - block_size, - ) + if self.fd_config.cache_config.enable_chunked_prefill or self.fd_config.cache_config.enable_prefix_caching: + latent_cache = forward_meta.caches[self.layer_id] if hasattr(forward_meta, "caches") else None + + block_size = self.mla_attn.block_size if hasattr(self.mla_attn, "block_size") else 64 + full_compressed_kv, full_k_pe = fused_read_cache_and_interleave( + latent_cache, + forward_meta.block_tables, + compressed_kv, + key_pe.squeeze(1), + forward_meta.cu_seqlens_k, + forward_meta.cu_seqlens_q, + self.kv_lora_rank, + self.qk_rope_head_dim, + block_size, + ) # Project latent KV to full key and value key_value = self.kv_b_proj(full_compressed_kv) diff --git a/fastdeploy/model_executor/pre_and_post_process.py b/fastdeploy/model_executor/pre_and_post_process.py index 3bf251009b3..50f764caa43 100644 --- a/fastdeploy/model_executor/pre_and_post_process.py +++ b/fastdeploy/model_executor/pre_and_post_process.py @@ -176,7 +176,7 @@ def pre_process( if specific_platform and not speculative_decoding: # Note(ZKK): This case's code is very simple! ids_remove_padding, batch_id_per_token, cu_seqlens_q, cu_seqlens_k = get_padding_offset( - input_ids, seq_lens_this_time, None, None, token_num_cpu + input_ids, seq_lens_this_time, None, None, seq_lens_decoder, token_num_cpu ) return ( ids_remove_padding, From 36b4339cb272d61bb2821c816f37a5968b7bcd32 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Fri, 8 May 2026 23:56:17 +0800 Subject: [PATCH 07/10] update tests --- .../test_mla_fused_read_interleave.py | 340 ++++++++++++++++++ 1 file changed, 340 insertions(+) create mode 100644 tests/model_executor/test_mla_fused_read_interleave.py diff --git a/tests/model_executor/test_mla_fused_read_interleave.py b/tests/model_executor/test_mla_fused_read_interleave.py new file mode 100644 index 00000000000..d1610bd8a8c --- /dev/null +++ b/tests/model_executor/test_mla_fused_read_interleave.py @@ -0,0 +1,340 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Precision parity tests for MLA fused read+interleave. + +Verifies that the Triton single-kernel implementation +``fused_read_cache_and_interleave_triton`` produces bit-level identical +results to the Python/Paddle reference ``fused_read_cache_and_interleave_naive`` +across a variety of batch/sequence configurations. +""" + +import numpy as np +import paddle +import pytest + +from fastdeploy.model_executor.layers.attention.mla_attention_backend import ( + fused_read_cache_and_interleave_naive, + fused_read_cache_and_interleave_triton, +) + +# Typical DeepSeek V3 MLA geometry. +KV_LORA_RANK = 512 +QK_ROPE_HEAD_DIM = 64 +LATENT_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM +BLOCK_SIZE = 64 + +pytestmark = pytest.mark.skipif( + not paddle.is_compiled_with_cuda() or paddle.device.cuda.device_count() == 0, + reason="fused_read_cache_and_interleave_triton requires CUDA", +) + + +def _build_inputs(cached_lens, new_lens, dtype="float16", seed=0): + """Construct all tensors needed by the fused read+interleave entry.""" + assert len(cached_lens) == len(new_lens) + bsz = len(cached_lens) + total_new = int(sum(new_lens)) + + # Paged latent cache with enough blocks for all batches. + max_blocks_per_seq = max(1, max((c + BLOCK_SIZE - 1) // BLOCK_SIZE for c in cached_lens) if bsz else 1) + # Give extra blocks so block ids are not all sequential. + num_blocks = max(bsz * max_blocks_per_seq + 7, 8) + + rng = np.random.default_rng(seed) + latent_np = rng.standard_normal((num_blocks, 1, BLOCK_SIZE, LATENT_DIM)).astype( + np.float16 if dtype == "float16" else np.float32 + ) + latent_cache = paddle.to_tensor(latent_np).cast(dtype) + + # block_tables: assign a distinct physical block id per (batch, block_idx). + bt_np = np.zeros((bsz, max_blocks_per_seq), dtype=np.int32) + free_ids = list(range(num_blocks)) + rng.shuffle(free_ids) + cursor = 0 + for b in range(bsz): + nb = (cached_lens[b] + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(nb): + bt_np[b, i] = free_ids[cursor] + cursor += 1 + block_tables = paddle.to_tensor(bt_np) + + new_kv_c_np = rng.standard_normal((max(total_new, 1), KV_LORA_RANK)).astype(np.float32) + new_k_pe_np = rng.standard_normal((max(total_new, 1), QK_ROPE_HEAD_DIM)).astype(np.float32) + new_compressed_kv = paddle.to_tensor(new_kv_c_np[:total_new] if total_new > 0 else new_kv_c_np[:0]).cast(dtype) + new_k_pe = paddle.to_tensor(new_k_pe_np[:total_new] if total_new > 0 else new_k_pe_np[:0]).cast(dtype) + + cu_total = np.zeros(bsz + 1, dtype=np.int32) + cu_new = np.zeros(bsz + 1, dtype=np.int32) + for i in range(bsz): + cu_total[i + 1] = cu_total[i] + cached_lens[i] + new_lens[i] + cu_new[i + 1] = cu_new[i] + new_lens[i] + cu_seqlens_k = paddle.to_tensor(cu_total) + cu_seqlens_q = paddle.to_tensor(cu_new) + + return { + "latent_cache": latent_cache, + "block_tables": block_tables, + "new_compressed_kv": new_compressed_kv, + "new_k_pe": new_k_pe, + "cu_seqlens_k": cu_seqlens_k, + "cu_seqlens_q": cu_seqlens_q, + } + + +def _run_both(inputs): + common = dict( + latent_cache=inputs["latent_cache"], + block_tables=inputs["block_tables"], + new_compressed_kv=inputs["new_compressed_kv"], + new_k_pe=inputs["new_k_pe"], + cu_seqlens_k=inputs["cu_seqlens_k"], + cu_seqlens_q=inputs["cu_seqlens_q"], + kv_lora_rank=KV_LORA_RANK, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + block_size=BLOCK_SIZE, + ) + + ref_kv, ref_pe = fused_read_cache_and_interleave_naive(**common) + out_kv, out_pe = fused_read_cache_and_interleave_triton(**common) + return ref_kv, ref_pe, out_kv, out_pe + + +def _assert_bitwise_equal(ref_kv, ref_pe, out_kv, out_pe): + assert tuple(out_kv.shape) == tuple(ref_kv.shape) + assert tuple(out_pe.shape) == tuple(ref_pe.shape) + # Both paths do pure copies (no arithmetic), so they should agree bit-for-bit. + np.testing.assert_array_equal(out_kv.cast("float32").numpy(), ref_kv.cast("float32").numpy()) + np.testing.assert_array_equal(out_pe.cast("float32").numpy(), ref_pe.cast("float32").numpy()) + + +# ----------------------------------------------------------------------------- +# Test cases: (name, cached_lens, new_lens) +# ----------------------------------------------------------------------------- +_CASES = [ + # Single batch, purely new tokens (no prefix cache). + ("single_no_cache", [0], [17]), + # Single batch, purely cached (edge case; triton path should still be correct). + ("single_all_cache_short", [13], [1]), + # Single batch spanning multiple cache blocks. + ("single_multi_block", [BLOCK_SIZE * 3 + 5], [9]), + # Single batch, cached aligned to exact block boundary. + ("single_cache_aligned", [BLOCK_SIZE * 2], [7]), + # Multi batch, uniform small lengths. + ("multi_uniform_small", [8, 8, 8, 8], [4, 4, 4, 4]), + # Multi batch, mixed: some with cache, some without. + ("multi_mixed_cache", [0, 33, 0, BLOCK_SIZE + 1], [5, 7, 11, 3]), + # Multi batch, long + short mixed. + ("multi_long_short", [BLOCK_SIZE * 5 + 3, 2, BLOCK_SIZE * 2, 0], [16, 1, 8, 12]), + # Larger batch with varied lengths. + ( + "multi_large", + [0, 17, BLOCK_SIZE, BLOCK_SIZE * 2 + 1, 5, 0, 64, BLOCK_SIZE * 4 + 30], + [8, 3, 16, 1, 12, 4, 9, 7], + ), + # Every cached_len > 0 (no batch without prefix cache). + ("multi_all_have_cache", [BLOCK_SIZE + 2, 7, BLOCK_SIZE * 2 + 17, 40], [6, 6, 6, 6]), + # bsz=1, minimal. + ("minimal", [1], [1]), +] + + +@pytest.mark.parametrize("name,cached_lens,new_lens", _CASES, ids=[c[0] for c in _CASES]) +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +def test_triton_matches_naive(name, cached_lens, new_lens, dtype): + inputs = _build_inputs(cached_lens, new_lens, dtype=dtype, seed=hash(name) & 0xFFFF) + ref_kv, ref_pe, out_kv, out_pe = _run_both(inputs) + _assert_bitwise_equal(ref_kv, ref_pe, out_kv, out_pe) + + +def test_total_tokens_zero(): + """All batches empty -> early return with empty tensors from both paths.""" + inputs = _build_inputs([0, 0], [0, 0], dtype="float16", seed=1) + ref_kv, ref_pe, out_kv, out_pe = _run_both(inputs) + assert ref_kv.shape[0] == 0 and out_kv.shape[0] == 0 + assert ref_pe.shape[0] == 0 and out_pe.shape[0] == 0 + + +# --------------------------------------------------------------------------- +# Hand-crafted baseline test: small tensors with deterministic integer values so +# every expected output entry can be written out and compared directly. This +# gives a ground-truth anchor independent of the naive reference, covering the +# real chunk-prefill / prefix-cache layout end-to-end. +# +# Scenario: +# bsz = 2 +# batch 0: cached=3, new=2 (spans 2 paged blocks: pbid 5 then pbid 2) +# batch 1: cached=0, new=4 +# block_size=2, kv_lora_rank=4, qk_rope_head_dim=2 +# +# Interleaved output layout (per-batch [cached, new]): +# t0..t2 : batch0 cached 0,1,2 (block_tables[0]=[5,2]) +# t3..t4 : batch0 new 0,1 +# t5..t8 : batch1 new 2,3,4,5 +# --------------------------------------------------------------------------- + + +def _fill_latent_cache(num_blocks, block_size, lora, rope): + """latent_cache[pbid, 0, off, d] = pbid * 1000 + off * 100 + d.""" + latent_dim = lora + rope + arr = np.zeros((num_blocks, 1, block_size, latent_dim), dtype=np.float32) + for p in range(num_blocks): + for o in range(block_size): + for d in range(latent_dim): + arr[p, 0, o, d] = p * 1000 + o * 100 + d + return arr + + +def test_manual_baseline(): + lora, rope = 4, 2 + block_size = 2 + num_blocks = 8 + + latent_np = _fill_latent_cache(num_blocks, block_size, lora, rope) + latent_cache = paddle.to_tensor(latent_np).cast("float32") + + # batch 0 cached tokens live in physical blocks (pbid=5 then pbid=2), + # batch 1 has no cached tokens so its row is unused. + block_tables = paddle.to_tensor(np.array([[5, 2], [0, 0]], dtype=np.int32)) + + # New tokens: 2 (batch 0) + 4 (batch 1) = 6 rows, deterministic values. + total_new = 6 + new_kv = np.zeros((total_new, lora), dtype=np.float32) + new_pe = np.zeros((total_new, rope), dtype=np.float32) + for i in range(total_new): + for d in range(lora): + new_kv[i, d] = -(i * 10 + d + 1) # negatives so they never clash with cache pattern + for d in range(rope): + new_pe[i, d] = -(i * 10 + d + 101) + new_compressed_kv = paddle.to_tensor(new_kv) + new_k_pe = paddle.to_tensor(new_pe) + + # cu_seqlens_k = cumsum(cached + new) per batch = [0, 5, 9] + # cu_seqlens_q = cumsum(new) per batch = [0, 2, 6] + cu_k = paddle.to_tensor(np.array([0, 5, 9], dtype=np.int32)) + cu_q = paddle.to_tensor(np.array([0, 2, 6], dtype=np.int32)) + + out_kv, out_pe = fused_read_cache_and_interleave_triton( + latent_cache, + block_tables, + new_compressed_kv, + new_k_pe, + cu_k, + cu_q, + lora, + rope, + block_size, + ) + + # Hand-built ground truth: total_tokens=9. + expected_kv = np.zeros((9, lora), dtype=np.float32) + expected_pe = np.zeros((9, rope), dtype=np.float32) + + # t0: batch0 cached 0 -> pbid=5, off=0 -> 5000 + 0 + d + # t1: batch0 cached 1 -> pbid=5, off=1 -> 5000 + 100 + d + # t2: batch0 cached 2 -> pbid=2, off=0 -> 2000 + 0 + d + cached_refs = [(5, 0), (5, 1), (2, 0)] + for t, (pbid, off) in enumerate(cached_refs): + for d in range(lora): + expected_kv[t, d] = pbid * 1000 + off * 100 + d + for d in range(rope): + expected_pe[t, d] = pbid * 1000 + off * 100 + lora + d + + # t3,t4: batch0 new tokens (src=0,1) + # t5..t8: batch1 new tokens (src=2,3,4,5) + new_sources = [0, 1, 2, 3, 4, 5] + for slot, src in enumerate(new_sources): + expected_kv[3 + slot] = new_kv[src] + expected_pe[3 + slot] = new_pe[src] + + np.testing.assert_array_equal(out_kv.numpy(), expected_kv) + np.testing.assert_array_equal(out_pe.numpy(), expected_pe) + + +def test_manual_baseline_no_cache(): + """All batches pure prefill (no prefix cache) -> kernel must still copy + new tokens in order with no off-by-one.""" + lora, rope = 4, 2 + block_size = 2 + # Any latent_cache / block_tables; kernel should never read from them. + latent_cache = paddle.zeros([4, 1, block_size, lora + rope], dtype="float32") + block_tables = paddle.zeros([3, 1], dtype="int32") + + new_kv = np.arange(5 * lora, dtype=np.float32).reshape(5, lora) + new_pe = np.arange(5 * rope, dtype=np.float32).reshape(5, rope) + 1000 + new_compressed_kv = paddle.to_tensor(new_kv) + new_k_pe = paddle.to_tensor(new_pe) + + # 3 batches, new=[2,1,2], no cache + cu_k = paddle.to_tensor(np.array([0, 2, 3, 5], dtype=np.int32)) + cu_q = paddle.to_tensor(np.array([0, 2, 3, 5], dtype=np.int32)) + + out_kv, out_pe = fused_read_cache_and_interleave_triton( + latent_cache, + block_tables, + new_compressed_kv, + new_k_pe, + cu_k, + cu_q, + lora, + rope, + block_size, + ) + np.testing.assert_array_equal(out_kv.numpy(), new_kv) + np.testing.assert_array_equal(out_pe.numpy(), new_pe) + + +def test_manual_baseline_all_cached(): + """Single batch, all cached (decode-like K build with zero new tokens).""" + lora, rope = 4, 2 + block_size = 2 + num_blocks = 4 + latent_np = _fill_latent_cache(num_blocks, block_size, lora, rope) + latent_cache = paddle.to_tensor(latent_np) + + # batch 0 cached=4 spans two blocks; pick pbid=3 then pbid=1. + block_tables = paddle.to_tensor(np.array([[3, 1]], dtype=np.int32)) + + # No new tokens; still need a valid (0-row) tensor. + new_compressed_kv = paddle.zeros([0, lora], dtype="float32") + new_k_pe = paddle.zeros([0, rope], dtype="float32") + + cu_k = paddle.to_tensor(np.array([0, 4], dtype=np.int32)) + cu_q = paddle.to_tensor(np.array([0, 0], dtype=np.int32)) + + out_kv, out_pe = fused_read_cache_and_interleave_triton( + latent_cache, + block_tables, + new_compressed_kv, + new_k_pe, + cu_k, + cu_q, + lora, + rope, + block_size, + ) + + # expected: pbid=3 off=0,1 then pbid=1 off=0,1 + expected_kv = np.zeros((4, lora), dtype=np.float32) + expected_pe = np.zeros((4, rope), dtype=np.float32) + for t, (pbid, off) in enumerate([(3, 0), (3, 1), (1, 0), (1, 1)]): + for d in range(lora): + expected_kv[t, d] = pbid * 1000 + off * 100 + d + for d in range(rope): + expected_pe[t, d] = pbid * 1000 + off * 100 + lora + d + np.testing.assert_array_equal(out_kv.numpy(), expected_kv) + np.testing.assert_array_equal(out_pe.numpy(), expected_pe) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 0c66a880262fbfd2a4ffe0692b964882a751f239 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Sat, 9 May 2026 00:25:59 +0800 Subject: [PATCH 08/10] update test --- tests/layers/test_attention_layer.py | 2 +- tests/operators/test_get_padding_offset.py | 4 +++- .../test_get_position_ids_and_mask_encoder_batch.py | 2 +- tests/operators/test_reasoning_phase_token_constraint.py | 6 ++++-- 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/layers/test_attention_layer.py b/tests/layers/test_attention_layer.py index 21d3deb5cff..56c6d6afc03 100644 --- a/tests/layers/test_attention_layer.py +++ b/tests/layers/test_attention_layer.py @@ -239,7 +239,7 @@ def create_forward_meta( input_ids = paddle.zeros([batch_size, max_model_len], dtype="int64") token_num = np.sum(seq_lens_this_time) ids_remove_padding, batch_id_per_token, cu_seqlens_q, cu_seqlens_k = get_padding_offset( - input_ids, seq_lens_this_time, None, None, token_num + input_ids, seq_lens_this_time, None, None, None, int(token_num) ) forward_meta = ForwardMeta( diff --git a/tests/operators/test_get_padding_offset.py b/tests/operators/test_get_padding_offset.py index 13c86c422d2..8d21660e89a 100644 --- a/tests/operators/test_get_padding_offset.py +++ b/tests/operators/test_get_padding_offset.py @@ -32,7 +32,9 @@ def test_get_padding_offset(self): batch_id_per_token, cu_seqlens_q, cu_seqlens_k, - ) = get_padding_offset(paddle.to_tensor(input_ids), paddle.to_tensor(seq_lens), None, None, token_num_cpu) + ) = get_padding_offset( + paddle.to_tensor(input_ids), paddle.to_tensor(seq_lens), None, None, None, int(token_num_cpu) + ) ref_x_remove_padding = np.array([8, 7, 8, 2, 4, 5, 5, 7, 6, 1, 7, 2, 6], "int64") ref_batch_id_per_token = np.array([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2, 2], "int32") diff --git a/tests/operators/test_get_position_ids_and_mask_encoder_batch.py b/tests/operators/test_get_position_ids_and_mask_encoder_batch.py index 2d1dd8e2f7c..cb7e38633fb 100644 --- a/tests/operators/test_get_position_ids_and_mask_encoder_batch.py +++ b/tests/operators/test_get_position_ids_and_mask_encoder_batch.py @@ -37,7 +37,7 @@ def test_basic_functionality(self): # Call the custom operator get_position_ids_and_mask_encoder_batch(seq_lens_encoder, seq_lens_decoder, seq_lens_this_time, position_ids) - expected_position_ids = np.array([0, 1, 2, 1, 0, 1, 2, 3], dtype=np.int32) + expected_position_ids = np.array([1, 2, 3, 2, 3, 0, 0, 0], dtype=np.int32) # Convert to numpy for comparison position_ids_np = position_ids.numpy() diff --git a/tests/operators/test_reasoning_phase_token_constraint.py b/tests/operators/test_reasoning_phase_token_constraint.py index 2c46b5d9c29..07fbf6e23ea 100644 --- a/tests/operators/test_reasoning_phase_token_constraint.py +++ b/tests/operators/test_reasoning_phase_token_constraint.py @@ -87,7 +87,8 @@ def setUp(self): seq_lens_output, None, None, - output_token_num.item(), + None, + int(output_token_num.item()), ) # self.output_padding_offset = paddle.zeros([self.token_num], dtype="int32") @@ -467,7 +468,8 @@ def test_perf_bsz128_vocab100k_status2(self): seq_lens_output, None, None, - output_token_num.item(), + None, + int(output_token_num.item()), ) # ------------------------ From 640ae84a5346881d23a08037fa2c5ee7038593f9 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Sat, 9 May 2026 11:20:24 +0800 Subject: [PATCH 09/10] optimize triton kernel --- .../layers/attention/mla_attention_backend.py | 122 +++++++++++------- 1 file changed, 76 insertions(+), 46 deletions(-) diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index ede903f32f0..b57618793b7 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -132,7 +132,6 @@ def fused_read_cache_and_interleave_naive( return full_compressed_kv, full_k_pe -@enable_compat_on_triton_kernel @triton.jit() def _fused_read_interleave_kernel( latent_cache_ptr, # [num_blocks, 1, block_size, LATENT_DIM] @@ -143,60 +142,75 @@ def _fused_read_interleave_kernel( block_tables_ptr, # [bsz, max_blocks_per_seq] int32 out_kv_c_ptr, # [total_tokens, kv_lora_rank] out_k_pe_ptr, # [total_tokens, qk_rope_head_dim] + total_tokens, bsz, max_blocks_per_seq: tl.constexpr, block_size: tl.constexpr, kv_lora_rank: tl.constexpr, qk_rope_head_dim: tl.constexpr, LATENT_DIM: tl.constexpr, + LOG2_MAX_BSZ: tl.constexpr, + BLOCK_M: tl.constexpr, ): - """Single-kernel fused read+interleave. - - Output tokens are laid out contiguously in batch-major order with per-batch - layout ``[cached, new]``. The owning batch ``b`` for each ``token_idx`` is - located via a linear scan on ``cu_total`` (bsz is small; O(bsz) overhead is - negligible next to per-token LATENT_DIM load/store). Per-batch cached - length is derived in-kernel as - ``nc = (cu_total[b+1] - cu_total[b]) - (cu_new[b+1] - cu_new[b])``. - """ - token_idx = tl.program_id(axis=0) - - # Linear scan: largest b such that cu_total[b] <= token_idx. - b = 0 - for i in range(bsz): - base_i = tl.load(cu_total_ptr + i) - if token_idx >= base_i: - b = i - - ct_b = tl.load(cu_total_ptr + b) - ct_b1 = tl.load(cu_total_ptr + b + 1) - cn_b = tl.load(cu_new_ptr + b) - cn_b1 = tl.load(cu_new_ptr + b + 1) + """Tiled fused read+interleave kernel. - k_len = ct_b1 - ct_b - nn = cn_b1 - cn_b - nc = k_len - nn - local_t = token_idx - ct_b + Each program handles ``BLOCK_M`` contiguous output tokens. The owning + batch of each token is recovered via an in-kernel binary search on the + cumsum ``cu_total`` (O(log bsz), fully L1-resident because bsz+1 is + small), avoiding a host-side ``repeat_interleave`` of shape + ``[total_tokens]``. + kv_c / k_pe are loaded as two separate vectors because ``kv_lora_rank`` + and ``qk_rope_head_dim`` are individually power-of-2 but their sum + (``LATENT_DIM``) is generally not (e.g. 512+64=576). + """ + pid = tl.program_id(axis=0) kv_c_offs = tl.arange(0, kv_lora_rank) k_pe_offs = tl.arange(0, qk_rope_head_dim) - if local_t < nc: - # cached token: walk paged latent_cache via block_tables - block_idx = local_t // block_size - block_off = local_t % block_size - pbid = tl.load(block_tables_ptr + b * max_blocks_per_seq + block_idx) - base = latent_cache_ptr + (pbid * block_size + block_off) * LATENT_DIM - kv_c_val = tl.load(base + kv_c_offs) - k_pe_val = tl.load(base + kv_lora_rank + k_pe_offs) - else: - # new token: linear region in new_compressed_kv / new_k_pe - src = cn_b + (local_t - nc) - kv_c_val = tl.load(new_kv_c_ptr + src * kv_lora_rank + kv_c_offs) - k_pe_val = tl.load(new_k_pe_ptr + src * qk_rope_head_dim + k_pe_offs) + # Unrolled at compile time - BLOCK_M is small + for m in tl.static_range(0, BLOCK_M): + token_idx = pid * BLOCK_M + m + if token_idx < total_tokens: + # Binary search: find largest b such that cu_total[b] <= token_idx. + # Loop bound LOG2_MAX_BSZ is compile-time; cu_total fits in L1. + lo = 0 + hi = bsz - 1 + for _ in tl.static_range(0, LOG2_MAX_BSZ): + mid = (lo + hi + 1) // 2 + cu_mid = tl.load(cu_total_ptr + mid) + if cu_mid <= token_idx: + lo = mid + else: + hi = mid - 1 + b = lo + + ct_b = tl.load(cu_total_ptr + b) + ct_b1 = tl.load(cu_total_ptr + b + 1) + cn_b = tl.load(cu_new_ptr + b) + cn_b1 = tl.load(cu_new_ptr + b + 1) + + k_len = ct_b1 - ct_b + nn = cn_b1 - cn_b + nc = k_len - nn + local_t = token_idx - ct_b + + if local_t < nc: + # cached token: walk paged latent_cache via block_tables + block_idx = local_t // block_size + block_off = local_t % block_size + pbid = tl.load(block_tables_ptr + b * max_blocks_per_seq + block_idx) + base = latent_cache_ptr + (pbid * block_size + block_off) * LATENT_DIM + kv_c_val = tl.load(base + kv_c_offs) + k_pe_val = tl.load(base + kv_lora_rank + k_pe_offs) + else: + # new token: linear region in new_compressed_kv / new_k_pe + src = cn_b + (local_t - nc) + kv_c_val = tl.load(new_kv_c_ptr + src * kv_lora_rank + kv_c_offs) + k_pe_val = tl.load(new_k_pe_ptr + src * qk_rope_head_dim + k_pe_offs) - tl.store(out_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) - tl.store(out_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) + tl.store(out_kv_c_ptr + token_idx * kv_lora_rank + kv_c_offs, kv_c_val) + tl.store(out_k_pe_ptr + token_idx * qk_rope_head_dim + k_pe_offs, k_pe_val) def fused_read_cache_and_interleave_triton( @@ -213,9 +227,13 @@ def fused_read_cache_and_interleave_triton( """Triton-accelerated fused read+interleave. Single kernel launch with zero Python metadata computation. All per-batch - geometry is derived in-kernel from ``cu_seqlens_k`` and - ``cu_seqlens_q``. Only one scalar D2H is needed: the final cumsum entry, - used to size the output tensors. + geometry is derived in-kernel from ``cu_seqlens_k`` and ``cu_seqlens_q`` + via a compile-time-bounded binary search. Only one scalar D2H is needed: + the final cumsum entry, used to size the output tensors. + + ``BLOCK_M=4`` and ``num_warps=8`` are tuned defaults (autotune is avoided + because ``triton.testing.do_bench`` is incompatible with paddle worker + subprocesses during ``profile_run``). """ bsz = cu_seqlens_k.shape[0] - 1 total_tokens = int(cu_seqlens_k[-1]) @@ -226,8 +244,16 @@ def fused_read_cache_and_interleave_triton( return full_compressed_kv, full_k_pe max_blocks_per_seq = block_tables.shape[1] + # Compile-time binary-search loop bound. ceil(log2(bsz)), clamped >=1. + log2_max_bsz = max(1, (bsz - 1).bit_length()) if bsz > 1 else 1 + + # Default tuned config: BLOCK_M=4 is robust across decode / prefill / + # chunk-prefill; num_warps=8 matches the 512-wide kv_lora_rank vector. + BLOCK_M = 4 + num_warps = 8 - _fused_read_interleave_kernel[(total_tokens,)]( + grid = ((total_tokens + BLOCK_M - 1) // BLOCK_M,) + _fused_read_interleave_kernel[grid]( latent_cache, new_compressed_kv, new_k_pe, @@ -236,12 +262,16 @@ def fused_read_cache_and_interleave_triton( block_tables, full_compressed_kv, full_k_pe, + total_tokens, bsz, max_blocks_per_seq=max_blocks_per_seq, block_size=block_size, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, LATENT_DIM=kv_lora_rank + qk_rope_head_dim, + LOG2_MAX_BSZ=log2_max_bsz, + BLOCK_M=BLOCK_M, + num_warps=num_warps, ) return full_compressed_kv, full_k_pe From 8159f572dad25c4f89d7812404cb6a7aa9bf7bf3 Mon Sep 17 00:00:00 2001 From: chang-wenbin Date: Sat, 9 May 2026 13:20:47 +0800 Subject: [PATCH 10/10] update deepseek-v3 modeling --- fastdeploy/model_executor/models/deepseek_v3.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index d75af311546..0ac87fa6dfa 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -209,6 +209,10 @@ def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None self.fd_config = fd_config self.layer_id = layer_id + self.block_size: int = fd_config.cache_config.block_size + self.enable_chunked_prefill = self.fd_config.cache_config.enable_chunked_prefill + self.enable_prefix_caching = self.fd_config.cache_config.enable_prefix_caching + self.use_gated_attn = getattr(self.fd_config.model_config, "use_gated_attn", False) self.use_bias = getattr(self.fd_config.model_config, "use_bias", False) self.tp_size = fd_config.parallel_config.tensor_parallel_size @@ -382,17 +386,15 @@ def forward( need_do_prefill = forward_meta.max_len_tensor_cpu[1] > 0 need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0 - if need_do_prefill: # max_enc_len_this_time + if need_do_prefill: # Handle prefix cache: read cached latent from paged cache and interleave # with the new-token latent in a single fused kernel call. full_compressed_kv = compressed_kv full_k_pe = key_pe.squeeze(1) - if self.fd_config.cache_config.enable_chunked_prefill or self.fd_config.cache_config.enable_prefix_caching: - latent_cache = forward_meta.caches[self.layer_id] if hasattr(forward_meta, "caches") else None + if self.enable_chunked_prefill or self.enable_prefix_caching: - block_size = self.mla_attn.block_size if hasattr(self.mla_attn, "block_size") else 64 full_compressed_kv, full_k_pe = fused_read_cache_and_interleave( - latent_cache, + forward_meta.caches[self.layer_id], forward_meta.block_tables, compressed_kv, key_pe.squeeze(1), @@ -400,7 +402,7 @@ def forward( forward_meta.cu_seqlens_q, self.kv_lora_rank, self.qk_rope_head_dim, - block_size, + self.block_size, ) # Project latent KV to full key and value