9 lines
709 B
Python
9 lines
709 B
Python
"""Patch sglang kimi-k3 image: fall back from tcgen05 fused-TMA attn_res on sm_120."""
|
|
p = "/sgl-workspace/sglang/python/sglang/srt/layers/attn_residual.py"
|
|
s = open(p).read()
|
|
old = " major, _ = torch.cuda.get_device_capability()\n _FAST_SUPPORTED = major >= 10"
|
|
new = " major, _ = torch.cuda.get_device_capability()\n # RTX 6000D is sm_120 (major 12): tcgen05 exists only on GB100/GB200/GB300\n # (majors 10/11), so sm_120 must take the Triton score/combine pipeline.\n _FAST_SUPPORTED = major in (10, 11)"
|
|
assert old in s, "pattern not found in attn_residual.py"
|
|
open(p, "w").write(s.replace(old, new))
|
|
print("OK patched attn_residual._use_fast -> majors (10,11)")
|