| """Fail-closed official FLA strict-config setup for one isolated process. |
| |
| No model/prompt/head changes. The guard prevents FLA STRICT's ordinary missing |
| key fallback and records actual configured calls. This is a diagnostic module, |
| not an installed change to the published wrapper or dependency environment. |
| """ |
| import hashlib,importlib,json,os,sys |
| from pathlib import Path |
|
|
| def sha(p):return hashlib.sha256(Path(p).read_bytes()).hexdigest() |
| def serialized(key):return json.dumps(key,separators=(',',':'),sort_keys=True) |
| def config_fields(config): |
| if isinstance(config,dict):return {k:config.get(k) for k in ['kwargs','num_warps','num_stages','num_ctas','maxnreg','ir_override']} |
| return {k:getattr(config,k,None) for k in ['kwargs','num_warps','num_stages','num_ctas','maxnreg','ir_override']} |
| def validate_profile(path,expected_sha): |
| path=Path(path).resolve() |
| if sha(path)!=expected_sha:raise ValueError('Profile hash changed') |
| profile=json.loads(path.read_text()) |
| if profile['format']!='decision-fla-l2norm-profile-v1' or profile.get('model_family')!='Qwen/Qwen3.5-4B' or profile['cache_mode']!='strict':raise ValueError('Profile format/mode unsupported') |
| if len(profile['files'])!=1 or profile['files'][0]['file']!='l2norm_fwd_kernel.json':raise ValueError('Unexpected profile file set') |
| f=path.parent/'l2norm_fwd_kernel.json' |
| if sha(f)!=profile['files'][0]['sha256']:raise ValueError('Explicit kernel config changed') |
| data=json.loads(f.read_text());entries={} |
| if data.get('default_config') is not None:raise ValueError('Implicit fallback defaults forbidden') |
| for h,item in data['autotune_entries'].items(): |
| key=item['autotune_key'];encoded=serialized(key) |
| if hashlib.md5(encoded.encode()).hexdigest()!=h or encoded in entries:raise ValueError('Invalid/duplicate key') |
| if len(key)!=5 or key[0]!=128 or type(key[1]) is not int or not 1<=key[1]<=64 or key[2:]!=['torch.bfloat16','torch.bfloat16','torch.float32']:raise ValueError('Unsupported numerical key') |
| c=item['config'] |
| if c['kwargs'].keys()!={'BT'} or c['kwargs']['BT'] not in [8,16,32,64] or c['num_warps'] not in [1,2,4,8,16] or c['num_stages']!=3 or c['num_ctas']!=1 or any(c.get(x) is not None for x in ['maxnreg','pre_hook','ir_override']):raise ValueError('Unexpected launch configuration') |
| entries[encoded]=c |
| if {json.loads(k)[1] for k in entries}!=set(range(1,65)):raise ValueError('Incomplete legal NB coverage') |
| return path,profile,entries |
|
|
| def attach_guard(kernel,cache_module,entries,telemetry): |
| original=kernel.run |
| def guarded(*args,**kwargs): |
| if cache_module.FLA_CACHE_MODE is not cache_module.FlaCacheMode.STRICT:raise RuntimeError('FLA strict mode changed') |
| key=cache_module.AutotuneKey.build(kernel.arg_names,kernel.keys,args,kwargs);encoded=serialized(list(key.autotune_key)) |
| if encoded not in entries:raise RuntimeError('Uncontracted FLA l2norm key: '+encoded) |
| expected=entries[encoded];loaded=cache_module.load_cached_config(kernel.kernel_name,key) |
| if config_fields(loaded)!=config_fields(expected):raise RuntimeError('FLA exact config lookup mismatch') |
| if key.autotune_key in kernel.cache and config_fields(kernel.cache[key.autotune_key])!=config_fields(expected):raise RuntimeError('A conflicting in-process kernel cache exists') |
| |
| |
| kernel.maybe_load_cached_config(key) |
| if key.autotune_key not in kernel.cache or config_fields(kernel.cache[key.autotune_key])!=config_fields(expected):raise RuntimeError('Official strict config did not load') |
| result=original(*args,**kwargs) |
| if config_fields(kernel.cache[key.autotune_key])!=config_fields(expected):raise RuntimeError('Kernel config changed during call') |
| telemetry['calls']+=1;telemetry['keys'][encoded]=telemetry['keys'].get(encoded,0)+1 |
| return result |
| kernel.run=guarded |
| return original |
|
|
| def install(profile_path,expected_sha): |
| path,profile,entries=validate_profile(profile_path,expected_sha) |
| if any(n=='fla' or n.startswith('fla.') for n in sys.modules):raise RuntimeError('Install profile before importing FLA; use a fresh isolated process') |
| for name,wanted in {'FLA_CACHE_MODE':'strict','FLA_CONFIG_DIR':str(path.parent)}.items(): |
| actual=os.environ.get(name) |
| if actual is not None and actual!=wanted:raise RuntimeError('Conflicting '+name) |
| os.environ[name]=wanted |
| import torch,triton,fla |
| actual={'torch':str(torch.__version__),'hip':torch.version.hip,'triton':triton.__version__,'fla':fla.__version__} |
| for name,value in actual.items(): |
| if value!=profile['runtime'][name]:raise RuntimeError('Runtime mismatch: '+name) |
| if not torch.cuda.is_available():raise RuntimeError('Profile is only qualified for the specified ROCm GPU') |
| arch=torch.cuda.get_device_properties(0).gcnArchName.split(':')[0] |
| if arch!=profile['runtime']['gpu_arch']:raise RuntimeError('Unsupported GPU architecture '+arch) |
| module=importlib.import_module('fla.modules.l2norm');cache_module=importlib.import_module('fla.ops.utils.cache');root=Path(fla.__file__).parent |
| for name,value in profile['fla_source_sha256'].items(): |
| if sha(root/name)!=value:raise RuntimeError('Pinned FLA source changed: '+name) |
| kernel=module.l2norm_fwd_kernel |
| if kernel.kernel_name!='l2norm_fwd_kernel' or kernel.keys!=['D','NB'] or kernel.cache:raise RuntimeError('Kernel identity or fresh-cache precondition failed') |
| telemetry={'profile_sha256':expected_sha,'status':'installed','calls':0,'keys':{},'strict_guard':True,'unknown_keys':'raise','autotune_fallback_permitted':False,'runtime':actual,'gpu_arch':arch,'process_scope':'one explicitly profiled Nox model; other model loading in this process is not supported'} |
| attach_guard(kernel,cache_module,entries,telemetry) |
| |
| |
| for name in ['l2norm_fwd_kernel1','l2norm_bwd_kernel','l2norm_bwd_kernel1']: |
| other=getattr(module,name) |
| def reject(*args,_name=name,**kwargs):raise RuntimeError('Uncontracted normalization kernel: '+_name) |
| other.run=reject |
| return telemetry |
|
|
