Uploaded sanitized BC250/ROCm Repository.
This commit is contained in:
Fabian
@@ -0,0 +1,156 @@
|
||||
"""Fix SDPA patch for BC-250 gfx1010: always use manual SDPA on CUDA"""
|
||||
import sys
|
||||
|
||||
path = '/home/fabian/ComfyUI/bc250_softmax_patch.py'
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
# Fix 1: patched_sdpa - ALWAYS use manual SDPA on CUDA (no threshold)
|
||||
old_sdpa = '''def patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
|
||||
S = key.size(-2)
|
||||
if query.is_cuda and S > SAFE_SOFTMAX_THRESHOLD:
|
||||
if not getattr(patched_sdpa, '_logged', False):
|
||||
logger.warning(f"[BC-250] Manual SDPA: Q={list(query.shape)}, S={S}")
|
||||
patched_sdpa._logged = True
|
||||
return _safe_sdpa(query, key, value, attn_mask=attn_mask,
|
||||
dropout_p=dropout_p, is_causal=is_causal, scale=scale)
|
||||
return _original_sdpa(query, key, value, attn_mask=attn_mask,
|
||||
dropout_p=dropout_p, is_causal=is_causal, scale=scale)'''
|
||||
|
||||
new_sdpa = '''def patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
|
||||
# gfx1010: ALWAYS use manual SDPA on CUDA — built-in math backend kernel hangs
|
||||
if query.is_cuda:
|
||||
if not getattr(patched_sdpa, '_logged', False):
|
||||
S = key.size(-2)
|
||||
logger.warning(f"[BC-250] Manual SDPA (ALWAYS): Q={list(query.shape)}, S={S}")
|
||||
patched_sdpa._logged = True
|
||||
return _safe_sdpa(query, key, value, attn_mask=attn_mask,
|
||||
dropout_p=dropout_p, is_causal=is_causal, scale=scale)
|
||||
return _original_sdpa(query, key, value, attn_mask=attn_mask,
|
||||
dropout_p=dropout_p, is_causal=is_causal, scale=scale)'''
|
||||
|
||||
if old_sdpa not in content:
|
||||
print("ERROR: patched_sdpa not found!")
|
||||
sys.exit(1)
|
||||
content = content.replace(old_sdpa, new_sdpa)
|
||||
print("OK: patched_sdpa → always manual on CUDA")
|
||||
|
||||
# Fix 2: _safe_sdpa - add sync after GPU matmul+softmax operations
|
||||
old_safe = '''def _safe_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
|
||||
L, S = query.size(-2), key.size(-2)
|
||||
if scale is None:
|
||||
scale = query.size(-1) ** -0.5
|
||||
attn_weight = torch.matmul(query, key.transpose(-2, -1)) * scale
|
||||
if is_causal:
|
||||
causal_mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=query.device), diagonal=1)
|
||||
attn_weight = attn_weight.masked_fill(causal_mask, float('-inf'))
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_weight = attn_weight.masked_fill(~attn_mask, float('-inf'))
|
||||
else:
|
||||
attn_weight = attn_weight + attn_mask
|
||||
attn_weight = _safe_softmax_impl(attn_weight, dim=-1)
|
||||
if dropout_p > 0.0:
|
||||
attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p)
|
||||
return torch.matmul(attn_weight, value)'''
|
||||
|
||||
new_safe = '''def _safe_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
|
||||
L, S = query.size(-2), key.size(-2)
|
||||
if scale is None:
|
||||
scale = query.size(-1) ** -0.5
|
||||
# gfx1010: sync between GPU ops to prevent kernel queue buildup + hang
|
||||
attn_weight = torch.matmul(query, key.transpose(-2, -1)) * scale
|
||||
if query.is_cuda:
|
||||
torch.cuda.synchronize()
|
||||
if is_causal:
|
||||
causal_mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=query.device), diagonal=1)
|
||||
attn_weight = attn_weight.masked_fill(causal_mask, float('-inf'))
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_weight = attn_weight.masked_fill(~attn_mask, float('-inf'))
|
||||
else:
|
||||
attn_weight = attn_weight + attn_mask
|
||||
attn_weight = _safe_softmax_impl(attn_weight, dim=-1)
|
||||
if query.is_cuda:
|
||||
torch.cuda.synchronize()
|
||||
if dropout_p > 0.0:
|
||||
attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p)
|
||||
output = torch.matmul(attn_weight, value)
|
||||
if query.is_cuda:
|
||||
torch.cuda.synchronize()
|
||||
return output'''
|
||||
|
||||
if old_safe not in content:
|
||||
print("ERROR: _safe_sdpa not found!")
|
||||
sys.exit(1)
|
||||
content = content.replace(old_safe, new_safe)
|
||||
print("OK: _safe_sdpa → sync after GPU ops")
|
||||
|
||||
# Fix 3: Also patch softmax to ALWAYS use manual on CUDA (same reason)
|
||||
old_softmax = '''def patched_softmax(input, dim=None, _stacklevel=3, dtype=None):
|
||||
if dim is None:
|
||||
dim = -1
|
||||
if dtype is not None:
|
||||
input = input.to(dtype)
|
||||
if input.is_cuda and input.shape[dim] > SAFE_SOFTMAX_THRESHOLD:
|
||||
if not getattr(patched_softmax, '_logged', False):
|
||||
logger.warning(f"[BC-250] Manual F.softmax: shape={list(input.shape)}, dim={dim}")
|
||||
patched_softmax._logged = True
|
||||
return _safe_softmax_impl(input, dim)
|
||||
return _original_softmax(input, dim=dim)'''
|
||||
|
||||
new_softmax = '''def patched_softmax(input, dim=None, _stacklevel=3, dtype=None):
|
||||
if dim is None:
|
||||
dim = -1
|
||||
if dtype is not None:
|
||||
input = input.to(dtype)
|
||||
# gfx1010: ALWAYS use manual softmax on CUDA — native kernel unreliable
|
||||
if input.is_cuda:
|
||||
if not getattr(patched_softmax, '_logged', False):
|
||||
logger.warning(f"[BC-250] Manual F.softmax (ALWAYS): shape={list(input.shape)}, dim={dim}")
|
||||
patched_softmax._logged = True
|
||||
return _safe_softmax_impl(input, dim)
|
||||
return _original_softmax(input, dim=dim)'''
|
||||
|
||||
if old_softmax not in content:
|
||||
print("WARNING: patched_softmax not found (may already be fixed)")
|
||||
else:
|
||||
content = content.replace(old_softmax, new_softmax)
|
||||
print("OK: patched_softmax → always manual on CUDA")
|
||||
|
||||
# Fix 4: Same for tensor.softmax
|
||||
old_tsm = '''def patched_tensor_softmax(self, dim=-1, dtype=None):
|
||||
if dtype is not None:
|
||||
self = self.to(dtype)
|
||||
if self.is_cuda and self.shape[dim] > SAFE_SOFTMAX_THRESHOLD:
|
||||
if not getattr(patched_tensor_softmax, '_logged', False):
|
||||
logger.warning(f"[BC-250] Manual softmax: shape={list(self.shape)}, dim={dim}")
|
||||
patched_tensor_softmax._logged = True
|
||||
return _safe_softmax_impl(self, dim)
|
||||
return _original_tensor_softmax(self, dim=dim)'''
|
||||
|
||||
new_tsm = '''def patched_tensor_softmax(self, dim=-1, dtype=None):
|
||||
if dtype is not None:
|
||||
self = self.to(dtype)
|
||||
# gfx1010: ALWAYS use manual softmax on CUDA
|
||||
if self.is_cuda:
|
||||
if not getattr(patched_tensor_softmax, '_logged', False):
|
||||
logger.warning(f"[BC-250] Manual softmax (ALWAYS): shape={list(self.shape)}, dim={dim}")
|
||||
patched_tensor_softmax._logged = True
|
||||
return _safe_softmax_impl(self, dim)
|
||||
return _original_tensor_softmax(self, dim=dim)'''
|
||||
|
||||
if old_tsm not in content:
|
||||
print("WARNING: patched_tensor_softmax not found (may already be fixed)")
|
||||
else:
|
||||
content = content.replace(old_tsm, new_tsm)
|
||||
print("OK: patched_tensor_softmax → always manual on CUDA")
|
||||
|
||||
# Fix 5: Update version string
|
||||
content = content.replace('v17 ready', 'v19 ready — ALL CUDA ops manual (no native kernels)')
|
||||
content = content.replace('Comprehensive Monkey-Patch v17', 'Comprehensive Monkey-Patch v19')
|
||||
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
|
||||
print("\nALL PATCHES APPLIED — v19")
|
||||
Reference in New Issue
Block a user