fix(OOM): skip flash_attn_varlen during profiling — use Q-tiling fallback
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:
@@ -200,6 +200,10 @@ FALLBACK_METHOD = '''
|
||||
max_seqlen = max(seq_lens_list)
|
||||
|
||||
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(
|
||||
q_flat.to(torch.float16),
|
||||
k_flat.to(torch.float16),
|
||||
|
||||
Reference in New Issue
Block a user