minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix FA3 prefill CP NaNs (#30439)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-07-07 20:16:09 -07:00
committed by GitHub
parent f3c3eea608
commit fa278a762c
2 changed files with 13 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashattention_backend.py
+11 -8
View File
@@ -45,6 +45,10 @@ from sglang.jit_kernel.flash_attention import (
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled 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 @triton.jit
def _build_pa_page_table_kernel( def _build_pa_page_table_kernel(
req_to_token_ptr, req_to_token_ptr,
@@ -361,14 +365,13 @@ class FlashAttentionBackend(AttentionBackend):
and not self.use_mla and not self.use_mla
) )
# Skip the FA3 scheduler_metadata precompute (PR #21104) under DP # Skip the FA3 scheduler_metadata precompute (PR #21104) when distributed
# attention. The precomputed buffer can become inconsistent with the # attention modes can change live cache_seqlens/num_splits across ranks.
# num_splits the C++ mha_fwd kernel derives from live cache_seqlens # A stale precomputed buffer can lead to an OOB read in the split-KV
# during decode, leading to an OOB read in the split-KV combine kernel # combine kernel (flash_fwd_combine_launch_template.h:52). Leaving
# (flash_fwd_combine_launch_template.h:52). Leaving scheduler_metadata # scheduler_metadata unset uses the existing per-layer metadata path.
# unset uses the existing per-layer metadata path. self._disable_scheduler_metadata_precompute = (
self._disable_scheduler_metadata_precompute = bool( _should_disable_scheduler_metadata_precompute(server_args)
getattr(server_args, "enable_dp_attention", False)
) )
def _compute_scheduler_metadata( def _compute_scheduler_metadata(
python/sglang/srt/model_executor/runner/eager_runner.py
+2
View File
@@ -305,6 +305,8 @@ class EagerRunner(BaseRunner):
) )
kwargs["input_embeds"] = sharded_hidden_states kwargs["input_embeds"] = sharded_hidden_states
forward_positions = sharded_positions forward_positions = sharded_positions
else:
forward_batch.attn_cp_metadata = None
category = ( category = (
"target_verify" "target_verify"