fix(OOM): return zeros during profiling — skip both flash_attn AND Q-tiling

flash_attn_varlen OOMs at 4096 tokens, Q-tiling also OOMs (K tensor too large).
During profiling (BI100_IN_STARTUP_PROFILE=1), return zeros immediately.
Profiling only measures memory footprint, not output correctness.

Restore: chunked_prefill=on, max_num_batched_tokens=4096.
This commit is contained in:
project6-dev
2026-08-13 16:36:17 +00:00
parent 048302bd4a
commit 20aac5b212
2 changed files with 6 additions and 1 deletions

View File

@@ -19,7 +19,7 @@ command:
- --disable-log-requests
- --disable-frontend-multiprocessing
- --max-num-batched-tokens
- '256'
- '4096'
- --enable-chunked-prefill
- --max-seq-len-to-capture
- '32768'

View File

@@ -219,6 +219,11 @@ FALLBACK_METHOD = '''
# Fallback: pure-math Q-tiling (original implementation)
_Q_CHUNK = 256
# During profiling, skip expensive attention — return zeros.
# Profiling only measures memory footprint, not output correctness.
if os.environ.get("BI100_IN_STARTUP_PROFILE") == "1":
return torch.zeros_like(query)
if (attn_metadata.query_start_loc is not None
and len(attn_metadata.query_start_loc) == num_seqs + 1):
q_lens = [