diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 85b6623ce..0f41886fc 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -45,6 +45,10 @@ from sglang.jit_kernel.flash_attention import ( from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled +def _should_disable_scheduler_metadata_precompute(server_args) -> bool: + return bool(server_args.enable_prefill_cp or server_args.enable_dp_attention) + + @triton.jit def _build_pa_page_table_kernel( req_to_token_ptr, @@ -361,14 +365,13 @@ class FlashAttentionBackend(AttentionBackend): and not self.use_mla ) - # Skip the FA3 scheduler_metadata precompute (PR #21104) under DP - # attention. The precomputed buffer can become inconsistent with the - # num_splits the C++ mha_fwd kernel derives from live cache_seqlens - # during decode, leading to an OOB read in the split-KV combine kernel - # (flash_fwd_combine_launch_template.h:52). Leaving scheduler_metadata - # unset uses the existing per-layer metadata path. - self._disable_scheduler_metadata_precompute = bool( - getattr(server_args, "enable_dp_attention", False) + # Skip the FA3 scheduler_metadata precompute (PR #21104) when distributed + # attention modes can change live cache_seqlens/num_splits across ranks. + # A stale precomputed buffer can lead to an OOB read in the split-KV + # combine kernel (flash_fwd_combine_launch_template.h:52). Leaving + # scheduler_metadata unset uses the existing per-layer metadata path. + self._disable_scheduler_metadata_precompute = ( + _should_disable_scheduler_metadata_precompute(server_args) ) def _compute_scheduler_metadata( diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 4d6f430e3..52a76fdb3 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -305,6 +305,8 @@ class EagerRunner(BaseRunner): ) kwargs["input_embeds"] = sharded_hidden_states forward_positions = sharded_positions + else: + forward_batch.attn_cp_metadata = None category = ( "target_verify"