From 7d2813cb51da612b47bf0d17a258b76f8c4140d8 Mon Sep 17 00:00:00 2001 From: Kevin Tu Date: Wed, 4 Feb 2026 14:52:56 -0600 Subject: [PATCH 1/3] Add support for compile only backend. Allows for compilation using AIU stack without AIUs present on system Signed-off-by: Kevin Tu --- aiu_fms_testing_utils/scripts/inference.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/aiu_fms_testing_utils/scripts/inference.py b/aiu_fms_testing_utils/scripts/inference.py index 45d52cc3..3076a167 100644 --- a/aiu_fms_testing_utils/scripts/inference.py +++ b/aiu_fms_testing_utils/scripts/inference.py @@ -138,6 +138,13 @@ default="inductor", choices=["inductor", "eager", "aot_eager"], ) +parser.add_argument( + "--compile_backend_aiu", + type=str, + help="Backend for AIU compilation", + default="sendnn", + choices=["sendnn", "sendnn_compile_only"], +) parser.add_argument( "--compile_dynamic", action="store_true", @@ -387,7 +394,10 @@ print("must set AIU_WORLD_RANK_0") exit() os.environ.setdefault("FLEX_COMPUTE", "SENTIENT") - os.environ.setdefault("FLEX_DEVICE", "PF") + if args.compile_backend_aiu == "compile_only_backend": + os.environ.setdefault("FLEX_DEVICE", "COMPILE") + else: + os.environ.setdefault("FLEX_DEVICE", "PF") device = torch.device("cpu") else: @@ -597,7 +607,7 @@ def select_int8_module( fx_config.backed_size_oblivious = "paged" in attn_name if is_aiu_backend: model.compile( - backend="sendnn", options={"sendnn.dynamic": args.compile_dynamic_sendnn} + backend=args.compile_backend_aiu, options={"sendnn.dynamic": args.compile_dynamic_sendnn} ) else: # compiling can make first inference pass slow From 86c8fa5ccced61e321d09a46593fe26d69ea89ba Mon Sep 17 00:00:00 2001 From: Kevin Tu Date: Wed, 4 Feb 2026 17:30:22 -0600 Subject: [PATCH 2/3] Fix a line with incorrect naming for "sendnn_compile_only" Signed-off-by: Kevin Tu --- aiu_fms_testing_utils/scripts/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiu_fms_testing_utils/scripts/inference.py b/aiu_fms_testing_utils/scripts/inference.py index 3076a167..b2c4a69f 100644 --- a/aiu_fms_testing_utils/scripts/inference.py +++ b/aiu_fms_testing_utils/scripts/inference.py @@ -394,7 +394,7 @@ print("must set AIU_WORLD_RANK_0") exit() os.environ.setdefault("FLEX_COMPUTE", "SENTIENT") - if args.compile_backend_aiu == "compile_only_backend": + if args.compile_backend_aiu == "sendnn_compile_only": os.environ.setdefault("FLEX_DEVICE", "COMPILE") else: os.environ.setdefault("FLEX_DEVICE", "PF") From 81ae14d0a667b6850860b7bfe2bd98c765b64b7f Mon Sep 17 00:00:00 2001 From: Kevin Tu Date: Wed, 25 Feb 2026 19:24:23 -0600 Subject: [PATCH 3/3] Fixed formatting issues Signed-off-by: Kevin Tu --- aiu_fms_testing_utils/scripts/inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/aiu_fms_testing_utils/scripts/inference.py b/aiu_fms_testing_utils/scripts/inference.py index b2c4a69f..502cc07f 100644 --- a/aiu_fms_testing_utils/scripts/inference.py +++ b/aiu_fms_testing_utils/scripts/inference.py @@ -607,7 +607,8 @@ def select_int8_module( fx_config.backed_size_oblivious = "paged" in attn_name if is_aiu_backend: model.compile( - backend=args.compile_backend_aiu, options={"sendnn.dynamic": args.compile_dynamic_sendnn} + backend=args.compile_backend_aiu, + options={"sendnn.dynamic": args.compile_dynamic_sendnn}, ) else: # compiling can make first inference pass slow