ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-04-20 23:30:16 +00:00
commit 019107ac9f
1 changed files with 19 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

29
scripts/benchmarks/qwen3_flex_inference.py
View file

@ -584,17 +584,23 @@ def main():
p.add_argument("--capture_cudagraph", action = "store_true")
p.add_argument("--lora_adapter", default = None)
# Kernel tuning (optional JSON-valued CLI args so we can sweep quickly):
p.add_argument("--decode_kernel_options", default = None,
help = "JSON for FlexKernelOptions applied in decode, "
"e.g. '{\"PRESCALE_QK\":true,\"USE_TMA\":true}'.")
p.add_argument("--prefill_kernel_options", default = None,
help = "Same but for prefill.")
p.add_argument(
"--decode_kernel_options",
default = None,
help = "JSON for FlexKernelOptions applied in decode, "
'e.g. \'{"PRESCALE_QK":true,"USE_TMA":true}\'.',
)
p.add_argument(
"--prefill_kernel_options", default = None, help = "Same but for prefill."
)
# If set, torch.compile the full attention-stack closure in addition to
# (or instead of) compiling just flex_attention. `reduce-overhead` is
# the interesting mode; it nests with our CUDA graph capture.
p.add_argument("--compile_model_forward", default = None,
choices = [None, "default", "reduce-overhead",
"max-autotune-no-cudagraphs"])
p.add_argument(
"--compile_model_forward",
default = None,
choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"],
)
p.add_argument("--stats_path", required = True)
args = p.parse_args()
@ -669,9 +675,12 @@ def main():
# function gets captured into the same graph.
if args.compile_model_forward:
torch._dynamo.config.cache_size_limit = 256
print(f"[flex] torch.compile(call_model_with_flex_kwargs, "
f"mode={args.compile_model_forward!r})")
print(
f"[flex] torch.compile(call_model_with_flex_kwargs, "
f"mode={args.compile_model_forward!r})"
)
import sys as _sys
_this = _sys.modules[__name__]
_this.call_model_with_flex_kwargs = torch.compile(
call_model_with_flex_kwargs,