diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index ab4cf841f..7680f2c6a 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -214,6 +214,16 @@ class FlashAttentionBackend(AttentionBackend): else 0 ) + # In embedding mode with no chunked prefill and radix cache disabled, + # skip KV cache write and use flash_attn_varlen_func with raw K/V + # instead of flash_attn_with_kvcache, bypassing paged KV cache entirely. + server_args = model_runner.server_args + self.fa_skip_kv_cache = ( + server_args.is_embedding + and server_args.chunked_prefill_size == -1 + and server_args.disable_radix_cache + ) + def _compute_scheduler_metadata( self, batch_size, max_seq_len_k, cache_seqlens, cu_seqlens_q ): @@ -606,7 +616,7 @@ class FlashAttentionBackend(AttentionBackend): and self.attn_cp_size > 1 ) - if save_kv_cache and not is_cp_mode: + if save_kv_cache and not is_cp_mode and not self.fa_skip_kv_cache: cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention @@ -764,6 +774,31 @@ class FlashAttentionBackend(AttentionBackend): self.device, _fa_cp_attn, ) + elif self.fa_skip_kv_cache: + # Embedding mode: skip KV cache read and use raw K/V tensors + # directly via flash_attn_varlen_func. The KV cache write is + # also skipped (guarded above). This eliminates store_kvcache + # and prepare_varlen_num_blocks overhead per layer. + assert k is not None, "fa_skip_kv_cache requires k to be provided" + assert k_descale is None and v_descale is None, ( + "fa_skip_kv_cache uses raw K/V tensors, " + "FP8 KV cache descaling is not supported in this mode" + ) + result = flash_attn_varlen_func( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k=k.view(-1, layer.tp_k_head_num, layer.head_dim), + v=v.view(-1, layer.tp_v_head_num, layer.v_head_dim), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_q, + softmax_scale=layer.scaling, + causal=causal, + window_size=window_size, + softcap=layer.logit_cap, + num_splits=self.num_splits, + **kwargs, + ) else: result = flash_attn_with_kvcache( q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),