ciflow/torchtitan/184918: [AOTI][5/N] Read Level 0 schema in save_gpu_kernel for AOTI (#184918)

Summary:

Modify save_gpu_kernel() to prefer the Level 0 launch_metadata_schema

dict (versioned, stable contract from 1/N) when available on

CompiledKernel.bin, falling back to hasattr probing of CompiledKernel

internals for older Triton versions.

Schema path reads: schema["entry_name"], schema["num_warps"],

schema["shared_mem"] — matching the abi_version=1 contract.

Also adds use_launch_metadata_schema config flag

(TORCHINDUCTOR_USE_LAUNCH_METADATA_SCHEMA env var, default True)

to control the new vs legacy metadata loading path.

Test Plan:

buck2 test fbcode//mode/opt --modifier ovr_config//triton:beta \

  fbcode//caffe2/test/inductor:test_triton_launcher_integration

Pass 6, Fail 0:

• SaveGpuKernelSchemaTest (4 unit tests, no GPU):
  
  
  
  • test_schema_path_reads_entry_name
  
  
  • test_schema_path_reads_num_warps
  
  
  • test_schema_path_reads_shared_mem
  
  
  • test_fallback_when_no_schema
  
  
  
  


• SaveGpuKernelAOTIE2ETest (2 E2E tests, H100):
  
  
  
  • test_aoti_pipeline_correctness: full AOTI export → compile → load → run
  
  
  • test_aoti_schema_populated: monkey-patch verifies launch_metadata_schema
    
    is populated and has expected keys (entry_name, num_warps, shared_mem)
  
  
  
  

Reviewed By: htyu

Differential Revision: D105888729