dylanyunlong/project_6_89d52222
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix(OOM): skip flash_attn_varlen during profiling — use Q-tiling fallback

Browse Source 
flash_attn_varlen_func allocates O(n²) temp memory for 4096 dummy tokens
during profile_run, causing OOM at gpu_memory_utilization=0.80.

BI100_IN_STARTUP_PROFILE=1 env var is already set by
patch_worker_startup_profile_guard.py during the synthetic forward pass.
Real inference requests still use flash_attn_varlen (much faster).
This commit is contained in:
project6-dev
2026-08-13 14:41:40 +00:00
parent 4518a39d12
commit bed1fc4d54
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
qwen3_6_scripts/patch_xformers_sdpa_seq.py
View File

@@ -200,6 +200,10 @@ FALLBACK_METHOD = '''
max_seqlen = max(seq_lens_list) max_seqlen = max(seq_lens_list)
try: try:
# Skip flash_attn during profiling — OOMs on large dummy batch
import os
if os.environ.get("BI100_IN_STARTUP_PROFILE") == "1":
raise RuntimeError("skip flash_attn during profiling")
out = _ixf.flash_attn_varlen_func( out = _ixf.flash_attn_varlen_func(
q_flat.to(torch.float16), q_flat.to(torch.float16),
k_flat.to(torch.float16), k_flat.to(torch.float16),