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

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

2
computility-run.yaml
View File

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

5
qwen3_6_scripts/patch_xformers_sdpa_seq.py
View File

@@ -219,6 +219,11 @@ FALLBACK_METHOD = '''
# Fallback: pure-math Q-tiling (original implementation) # Fallback: pure-math Q-tiling (original implementation)
_Q_CHUNK = 256 _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 if (attn_metadata.query_start_loc is not None
and len(attn_metadata.query_start_loc) == num_seqs + 1): and len(attn_metadata.query_start_loc) == num_seqs + 1):
q_lens = [ q_lens = [