diff --git a/.gitattributes b/.gitattributes
index d3846436c0cf0147f2cde6a78b764bfd7d6ab535..892021c092d92ebb47a57164376013df3551929c 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -15914,3 +15914,15 @@ neuronxcc-2.21.33363.0+82129205/MODULE_e32b45ec3e4f6da82d81+9c024cc2/model.neff
 neuronxcc-2.21.33363.0+82129205/MODULE_e32b45ec3e4f6da82d81+9c024cc2/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
 neuronxcc-2.21.33363.0+82129205/MODULE_f7c37b12c4228f956d12+3d105e60/model.neff filter=lfs diff=lfs merge=lfs -text
 neuronxcc-2.21.33363.0+82129205/MODULE_f7c37b12c4228f956d12+3d105e60/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.neff filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.neff filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.neff filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.neff filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.neff filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.neff filter=lfs diff=lfs merge=lfs -text
+neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/wrapped_neff.hlo filter=lfs diff=lfs merge=lfs -text
diff --git a/neuronxcc-2.21.33363.0+82129205/0_REGISTRY/0.4.5.dev1/llama/unsloth/Llama-3.2-1B-Instruct/f18d2067d0df4f48c0c4.json b/neuronxcc-2.21.33363.0+82129205/0_REGISTRY/0.4.5.dev1/llama/unsloth/Llama-3.2-1B-Instruct/f18d2067d0df4f48c0c4.json
new file mode 100644
index 0000000000000000000000000000000000000000..b09933cb934911dd7e5eaebb310543a8ac5ebe18
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/0_REGISTRY/0.4.5.dev1/llama/unsloth/Llama-3.2-1B-Instruct/f18d2067d0df4f48c0c4.json
@@ -0,0 +1,63 @@
+{
+  "_entry_class": "SingleModelCacheEntry",
+  "_model_id": "unsloth/Llama-3.2-1B-Instruct",
+  "_task": "text-generation",
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "dtype": "bfloat16",
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "neuron": {
+    "_serialized_key": "NxDNeuronConfig",
+    "batch_size": 1,
+    "capacity_factor": null,
+    "checkpoint_id": "unsloth/Llama-3.2-1B-Instruct",
+    "checkpoint_revision": null,
+    "continuous_batching": false,
+    "ep_degree": 1,
+    "fused_qkv": false,
+    "glu_mlp": true,
+    "local_ranks_size": 2,
+    "max_batch_size": 1,
+    "max_context_length": 4096,
+    "max_topk": 256,
+    "n_active_tokens": 4096,
+    "neuronxcc_version": "2.21.33363.0+82129205",
+    "on_device_sampling": false,
+    "optimum_neuron_version": "0.4.5.dev1",
+    "output_logits": false,
+    "pp_degree": 1,
+    "sequence_length": 4096,
+    "speculation_length": 5,
+    "start_rank_id": 0,
+    "target": "trn1",
+    "torch_dtype": "bfloat16",
+    "tp_degree": 2
+  },
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "llama3"
+  },
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": true,
+  "unsloth_fixed": true,
+  "use_cache": true,
+  "vocab_size": 128256
+}
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_10217061096959125489+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_10217061096959125489+e30acd3a/model.neff
index 396de1e26db1f044fe928bd3f841bb34c40bbf91..40b9f07b97e90b2cba87f9fba3e7d455369e7595 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_10217061096959125489+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_10217061096959125489+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_10244305442015770634+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_10244305442015770634+e30acd3a/model.neff
index 847c9dfe35c8bc70486b0d6ee6b85f848be95c70..233b166b0e26a8ea36fa1b3559a009004ca0f34c 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_10244305442015770634+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_10244305442015770634+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_10606948783918825529+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_10606948783918825529+e30acd3a/model.neff
index 8eb1747ca673d5aaa6147c1b997d5f112745f7b2..98f2dfa94a6a125b72978d5b23a8d2734baeef05 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_10606948783918825529+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_10606948783918825529+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_10645643398657092095+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_10645643398657092095+e30acd3a/model.neff
index 5edede1af6ca6c59590a6609d29978116cdf6e52..d6e7dd120ad37c442466a878f19e67d27442486a 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_10645643398657092095+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_10645643398657092095+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_10661660426924300837+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_10661660426924300837+e30acd3a/model.neff
index e35b4100a1b0a580fd9fcd4bd30a2c8216364769..cf7ee39fdf46ee8d324e2b74dd86fa8106e0cdc1 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_10661660426924300837+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_10661660426924300837+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_10746122569655005679+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_10746122569655005679+e30acd3a/model.neff
index 2346df04633103d3867d97a503a9b6edd68a9133..613caa4feb9fcc97dca9515b9034b1112ca683f0 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_10746122569655005679+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_10746122569655005679+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_11086318750207148626+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_11086318750207148626+e30acd3a/model.neff
index ad4bf1eb81db096412d165edb73cb7f066b13701..6b13e86c814a59a920db2f977faa7c9ce1d8bd90 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_11086318750207148626+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_11086318750207148626+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_11909310600244571805+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_11909310600244571805+e30acd3a/model.neff
index 3b7c49e2ccf7124fa63371e9aa18038356f1270e..411137fe5a48232b3e383458261a00b8408a23e1 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_11909310600244571805+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_11909310600244571805+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_11933818254123612383+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_11933818254123612383+e30acd3a/model.neff
index 5c27532819dc89b33d0f75bca27ad9437a48e45a..60cccacc5655775438d245198a9d29e04ff3c3a1 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_11933818254123612383+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_11933818254123612383+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1210392327607194823+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1210392327607194823+e30acd3a/model.neff
index 29fa6c963c2cb8ded9061d2685a9eae3de7af707..5b30e1137439a885cc33f44f6ffa061915eec00e 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1210392327607194823+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1210392327607194823+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_12487216553200321032+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_12487216553200321032+e30acd3a/model.neff
index 0fa12b9ce63f44ca251f957974b312faf6218d9d..37eb5433e477b7aa27ec7740b6132a1b8531b72f 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_12487216553200321032+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_12487216553200321032+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_13085549342645515693+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_13085549342645515693+e30acd3a/model.neff
index 88731010bd56a10c63ac2a9788381b8a5979274b..1435bdb4b44e9f24e17c7d26a2d554478bd509d9 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_13085549342645515693+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_13085549342645515693+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_13102010590082783346+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_13102010590082783346+e30acd3a/model.neff
index 125dc2121ea2e4095eb5eadc61f61f080b1064af..07c4907b48018d25f22fc79244e28d104b82198b 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_13102010590082783346+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_13102010590082783346+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_13475273426270778455+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_13475273426270778455+e30acd3a/model.neff
index 051c28e8d6902bedb90cf415b8549a41994aa933..8abe29d89c6df2111132dab6319464d18ae3914b 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_13475273426270778455+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_13475273426270778455+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_13728813963059599796+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_13728813963059599796+e30acd3a/model.neff
index 1de531d1dc23abb1ac50f875c881573dd6645100..3162e1787765753374d60230b259875927f53d11 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_13728813963059599796+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_13728813963059599796+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_13783872956654054643+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_13783872956654054643+e30acd3a/model.neff
index fa3810a93e3e29de3e448d9318894b098fc68a9b..8844d03f77bda17a9e0d1ed0fa4dbde3f82f8004 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_13783872956654054643+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_13783872956654054643+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_14097776130612710282+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_14097776130612710282+e30acd3a/model.neff
index 499a712853f81f267e4465b3064b92d5983c32cc..e8505996676d1fd7a06d16642643c4933ab48692 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_14097776130612710282+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_14097776130612710282+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_14155105943057125178+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_14155105943057125178+e30acd3a/model.neff
index 94e8ae73ea5ce5e1523e2e744137e05f46baad7e..4a9e61a241b97702fc348875f11c925843e5f23c 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_14155105943057125178+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_14155105943057125178+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_14173682194645188821+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_14173682194645188821+e30acd3a/model.neff
index 9f2ac3c55d59f149caf7eba2ce0d33a4bca71b2b..ada252ed1d77621fab7b434cb426878503ea5aea 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_14173682194645188821+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_14173682194645188821+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_14274833313744358348+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_14274833313744358348+e30acd3a/model.neff
index 9e937827865f39204c7275c8a20e09d9f029b3e7..0bbf091caa031ecc0f1acb5ca1da7840b99acd4e 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_14274833313744358348+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_14274833313744358348+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_14469716438804935215+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_14469716438804935215+e30acd3a/model.neff
index b87804c93e647a43ff275bd992b2eb7220c1e4ac..9533834641a892ae2b553c84fe5fdd9b3727c602 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_14469716438804935215+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_14469716438804935215+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_14739672283771660808+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_14739672283771660808+e30acd3a/model.neff
index ca9480b7b9a98d3c5e13244639627dba76c04267..a4ebcd273e5f83971a87f683b20b9f013634a511 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_14739672283771660808+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_14739672283771660808+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15442663025941492357+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15442663025941492357+e30acd3a/model.neff
index 8292709ec6310315d049f763438eef5bc0dbe4e0..9e6aa3ce2dd4f436cd5c7690ab09d99370d63766 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15442663025941492357+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15442663025941492357+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15557971880365771457+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15557971880365771457+e30acd3a/model.neff
index d1ab6087decb3b62a9abf6c579c13d2c3d74fd82..ad1de7a93f7ac08ae4fffb3a99cea99e681a054e 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15557971880365771457+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15557971880365771457+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15564425602929126510+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15564425602929126510+e30acd3a/model.neff
index 9d01a80cdff9bda6947589f41c632c78c97f11be..ccdc5d40ef99d3d5c144ed456432a353dde78781 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15564425602929126510+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15564425602929126510+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15589838460896944293+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15589838460896944293+e30acd3a/model.neff
index b23d3ac65181f9235b3e9fce7dbac2afa6f39821..f84553859b5c87c971add434bc897f0a936b4ad8 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15589838460896944293+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15589838460896944293+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15654572849171857535+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15654572849171857535+e30acd3a/model.neff
index 2db752f101f81b8dbb74877158a7ee250530f526..7602fc1274b2fd34c575711a7bd223d7e0dea4f7 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15654572849171857535+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15654572849171857535+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15767383571209512795+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15767383571209512795+e30acd3a/model.neff
index 773cfd46dce4feeb803ba0ec0011a5b869c3a6df..ab32d86cc53181d4c74fd3bd64fd0081090feab4 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15767383571209512795+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15767383571209512795+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15888836045088309511+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15888836045088309511+e30acd3a/model.neff
index 00605ed4b9f24a1ddf8359c47f775c2950f5ac20..c6710e32e7d4afaf44a36e0e85fefe650c5771d2 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15888836045088309511+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15888836045088309511+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_15929451261464042997+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_15929451261464042997+e30acd3a/model.neff
index 9b95227b85668512abac8027e44b14f2d3bf8a8c..b018af6d3850ae4d0f398547e18372b06e3b6e2b 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_15929451261464042997+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_15929451261464042997+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_16016505958416521648+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_16016505958416521648+e30acd3a/model.neff
index d06adcd6defa03c9c975421db3839fc71b095b91..946b3f29d49547c27018805f291cd58f7a336356 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_16016505958416521648+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_16016505958416521648+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_16053163479112702088+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_16053163479112702088+e30acd3a/model.neff
index a9856281b165bfa8c04cc5e6c3bbe6a38d609751..983321b3e2297f4971b5fbe910177fdbcb4a3937 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_16053163479112702088+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_16053163479112702088+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_16204387852795926216+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_16204387852795926216+e30acd3a/model.neff
index c2caafb890e274aa2599d13abc4937c7dc8ca295..1679e1c3d794eaed7cb005a50dff4c167a6c5941 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_16204387852795926216+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_16204387852795926216+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1656719109221189948+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1656719109221189948+e30acd3a/model.neff
index 9c57b68dc4f96df85509d07b3c3cf9b012818857..8a250fe24158b5518b6f8df0debfdb852ff8c518 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1656719109221189948+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1656719109221189948+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_16599571375348449904+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_16599571375348449904+e30acd3a/model.neff
index 0c448b556bedaaa4c44da7c9bca266ece27bcd3e..871eb7efa51488cac61b3b066e114a4991e542d7 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_16599571375348449904+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_16599571375348449904+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1663757798483801648+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1663757798483801648+e30acd3a/model.neff
index 5d4b30457b02d35e24545636c96fc5f205ee658c..0852cb0b281a32a395512fb9e3337fa03bd9bfe5 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1663757798483801648+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1663757798483801648+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_16738296820980389103+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_16738296820980389103+e30acd3a/model.neff
index c3d51b252669148c4d8536c41d27f64b91cb3d69..1e9a9528c49f90c43202e2786008297c616625f7 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_16738296820980389103+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_16738296820980389103+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1696825468766062114+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1696825468766062114+e30acd3a/model.neff
index 20a1bb09593172aeab5d62117ecb665d25d714dd..1fc1eaffa2260a806c46616d334b7558eac44a3d 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1696825468766062114+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1696825468766062114+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_17476209562158013765+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_17476209562158013765+e30acd3a/model.neff
index dec12d7117742ee00ff3320565c72ca1efd8c879..80fa19ee5fdd162d4866086dcb15c4042d357de6 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_17476209562158013765+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_17476209562158013765+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_17781596253725927902+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_17781596253725927902+e30acd3a/model.neff
index 52915fa05955438ea603b320065b0c574b07f0bc..3df3cd67d43bfcf521206887f4951deeb0856acd 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_17781596253725927902+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_17781596253725927902+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1795740353831178306+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1795740353831178306+e30acd3a/model.neff
index b31b9bd39e56098e1d1716394c57eed4f54c0482..d412f45f2588692d97fadcf2e3f7ac1c1d713d74 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1795740353831178306+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1795740353831178306+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1800832390737682969+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1800832390737682969+e30acd3a/model.neff
index 3273cbfa589d45a18736df55991e08ebd98393c2..f169ab6bd00a36bd4c9126bca4f888d735d406fb 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1800832390737682969+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1800832390737682969+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1805481651134498710+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1805481651134498710+e30acd3a/model.neff
index 0dce33f33dc7183ce80a1be993325a833d000222..d6304ec4b6ed2b10b6da44b43b87307f6a6b13c4 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1805481651134498710+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1805481651134498710+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_1805728312057401221+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_1805728312057401221+e30acd3a/model.neff
index 80e2386ba424d5714ba41d91704479d37d2db1be..8e131ee5ca1170cd65930b96e2c1c1e3ea423d4d 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_1805728312057401221+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_1805728312057401221+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_18070749384531238695+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_18070749384531238695+e30acd3a/model.neff
index 0a9f157340dc3579fd26ea2589f4ae7ea0cbb05f..c1f6a85dc68f9d6f2f10d5cd8b1dcac5fbcc33cb 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_18070749384531238695+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_18070749384531238695+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_18134517905922687964+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_18134517905922687964+e30acd3a/model.neff
index dc96d04727757784bff5c0b6e198aa9ffc6ed1cf..98cfb8eef221e6b752c896cb1333796c8e757287 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_18134517905922687964+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_18134517905922687964+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_18207475553554060390+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_18207475553554060390+e30acd3a/model.neff
index 6f3f782296e35b5dc7c306ac22d799e028a988da..02e03403c5a84841bd1a29692231ea8b87de75cb 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_18207475553554060390+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_18207475553554060390+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2157613318347839507+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2157613318347839507+e30acd3a/model.neff
index b4bda30533d3bd170ce324688837121ba1376807..6897d20795d1b91ea15241c5aa9b3ee4b84522c6 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2157613318347839507+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2157613318347839507+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2203200257461345827+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2203200257461345827+e30acd3a/model.neff
index 8d01c36eed73220932ef38614d75c7f2c21c74ac..818bd9a35edb05f2e67f70208d0394557ad02116 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2203200257461345827+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2203200257461345827+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2410426589418994260+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2410426589418994260+e30acd3a/model.neff
index 15de94c558d864912a252c46cbd43dd256ea5adc..57719f65b79a1216d9e956665e88dbec703c161d 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2410426589418994260+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2410426589418994260+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2426414314187505427+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2426414314187505427+e30acd3a/model.neff
index 8b34f2f09a3312ebc9071c5e4ceed734bc7e2df7..c380b67658da14a40864698e0a45f73d8c939c0b 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2426414314187505427+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2426414314187505427+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2576759111807165188+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2576759111807165188+e30acd3a/model.neff
index 1dc4ae56ee0c5a03d2028d5626412ffce5af1ab8..916372c4c0ae25521b74a79dcdb8c8e8d648372e 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2576759111807165188+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2576759111807165188+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2650447136112456251+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2650447136112456251+e30acd3a/model.neff
index 05413e7711a52c65bb2b05f0f2a7f4568d3d7d84..3ea2786feb2d1d1d942f98e5460b14506b41c9ae 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2650447136112456251+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2650447136112456251+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2682953264353234433+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2682953264353234433+e30acd3a/model.neff
index a637e9ec1a397575e2fc9dab90fce9083baa3e31..1f8d36067d27336c49056f88579091b13bcbed60 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2682953264353234433+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2682953264353234433+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2699280712748688265+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2699280712748688265+e30acd3a/model.neff
index 1f557390076d0f892c0e0a6b170d7f625c9f0d7a..df69c695e37cd594377f113532574a5d00a5556f 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2699280712748688265+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2699280712748688265+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2811495285170804454+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2811495285170804454+e30acd3a/model.neff
index b9563e5eb35a0335b353de8544a7e118e873f1dc..e2aa78050d550752848a3615c44527f6be49e268 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_2811495285170804454+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_2811495285170804454+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.hlo_module.pb
index 5b31a795fa051fe6382001027af513aab21cff12..1b8352e1ee9da452fd713c1a756c86482df2737e 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.hlo_module.pb
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.hlo_module.pb
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a587d5ddb67a42533df39b1b05311d022b9cb500122dfc7b382e9f459a2b148
+oid sha256:e635e3fd90e62f2f22ffdf53f78d7bdd4fa50499755a35b844ef028afd5c49ab
 size 593552
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.neff
index 91ad71f9e9e6bfc86ba5afb61a45f08baa4415b1..889b497a78bc72ee37105be5599c09ac81925caa 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.neff
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_2ce3d3065088729c7d7f+24129607/model.neff
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1dd568378b23fc6e1c40644733724955b7d6c831df572a8ebe1a15fe02623f5d
+oid sha256:1749e3e46f0b969b54be0d5b6ee3d60675511f3c6c8daaa2ad28cc1ecfe9c47e
 size 3769344
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_302950992539679003+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_302950992539679003+e30acd3a/model.neff
index 33f956973b07b5b8376ce847dab11a393e2f3c06..f588109c098cb8a807858125c07c44a526fef9a3 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_302950992539679003+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_302950992539679003+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..b1192c97fb3b066cc611a8065f723e9e3d6cff37
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_8049efec-dc22-437f-b199-674be5a4eabe/compiler_workdir/Llama3RotaryEmbedding/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..f1e2473c54cffac949ecf31d3b6ec636c290545b
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4a29d4cf27aec9af90e53324eff9ee69e033679482f4016a9dd3e5416c7c68a
+size 1584
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..4633eb93c6b71f3d3090d1c70d2a36d53f38eedd
Binary files /dev/null and b/neuronxcc-2.21.33363.0+82129205/MODULE_3a1c18c55a510c478222+1c7e690c/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4013984518400530928+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4013984518400530928+e30acd3a/model.neff
index b97983648b2e5c658ba02c0808b76b8d478b6a9f..23b0e26710c2143186bed254297026d5526a9a54 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4013984518400530928+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4013984518400530928+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4052264821257342969+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4052264821257342969+e30acd3a/model.neff
index 5401c71e4537b02afb4ef778ea861bb1b93dfe7d..a6de587c5e8ec0f9583191da7c85df3fb93b6c2a 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4052264821257342969+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4052264821257342969+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4166695227062860792+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4166695227062860792+e30acd3a/model.neff
index 336c377813c8e58ed6bf7459121d00488533145e..1618edc292c8a5b7bb8f106e24a6d6660ab3d43f 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4166695227062860792+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4166695227062860792+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4232557591140652245+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4232557591140652245+e30acd3a/model.neff
index b331e65527e1f029ddeb95e6bbf91decf6417c53..3088e0da5c9f967a2655a170105e55d49dda2670 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4232557591140652245+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4232557591140652245+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4606491920565381523+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4606491920565381523+e30acd3a/model.neff
index 53acc2e64d4abea2dad0a6830d7d80f3ceb66c14..6a9805121916325204584c79bea850c9d5f95b9d 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4606491920565381523+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4606491920565381523+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4608333438381222049+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4608333438381222049+e30acd3a/model.neff
index ebc6b1c235e37111300c3e03183b4d1ee60ed4d7..e405934e9cb11825efbeb35904987933ca41f304 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4608333438381222049+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4608333438381222049+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4683016291147007110+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4683016291147007110+e30acd3a/model.neff
index ac0f83b0565fbf5eb25b2808dc4be02ef859442a..d4000cb2c6c1cd3fc18e3c2f813a341b49819f3f 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4683016291147007110+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4683016291147007110+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4816422433712790690+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4816422433712790690+e30acd3a/model.neff
index 1bb890344053af37dd8f61da3c9892e1f0d45087..96e9f2ae81358ac9ae8e03e03e12b8a77d90ead5 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4816422433712790690+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4816422433712790690+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4842952601010372434+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4842952601010372434+e30acd3a/model.neff
index 10f7c1a7a50f128fbd559f075fb88e65b6f53252..7e1332780fac48e35ddb585b362166dafe86d08c 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4842952601010372434+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4842952601010372434+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_4939721357859779936+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_4939721357859779936+e30acd3a/model.neff
index aafb2ae6080d8b6a6f444b7fddff8a9020246f02..a455378265f4f7120b7ef863b239f8fa6766b348 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_4939721357859779936+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_4939721357859779936+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_5397223456911199516+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_5397223456911199516+e30acd3a/model.neff
index ba07a99684c80c8533bfacacdcd9abaeda809d9b..ede5737b6f7ac1ee8066ab1c2584bc843a8a9f4a 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_5397223456911199516+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_5397223456911199516+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_5724544066416780383+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_5724544066416780383+e30acd3a/model.neff
index 2d2c94139a41194fb38e0d2e510fb7b1e59101b5..127d91398b239dbf35632f1be68af719df909927 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_5724544066416780383+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_5724544066416780383+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_5765916484866650909+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_5765916484866650909+e30acd3a/model.neff
index 43ed2b59bda90848866ce9ad6a8b86d45229ab3a..b1581f5b943eadad362a911a6e4fa40e4c3f64a7 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_5765916484866650909+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_5765916484866650909+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..2a8d53e583967f4a388fe6ff3d17dc4ccb45367b
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_475dd50d-500f-419e-a600-dd04cf0ce4df/compiler_workdir/AttentionModelWrapper/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..a11cc062167102b0f4faa9fdf16388d72d9e0015
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4590b82eb167fd46963abdbb717cc18cb9511961a3393d100eeacdbdc45b0243
+size 11280
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..369dceadeb31e92b72fdf9e85f0714f7871c7ab4
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/model.neff
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f1eac2be095048dab2eab7e10bf8126c0006602088bcaa29669b3cea780ce7e
+size 1444864
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/wrapped_neff.hlo
new file mode 100644
index 0000000000000000000000000000000000000000..ce0a9f283fb01e77e85a1aaa429f562996cf770b
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_62d33aaee969bb42c631+0382c9c8/wrapped_neff.hlo
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33c829ae3786c8c5170456793d7cf4042511b65b9b5516e5f772ec9c83d6b5f6
+size 1447845
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_6368409506294683105+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_6368409506294683105+e30acd3a/model.neff
index 8b7c46f71774982f6fb4c99d1975e1c3ffe5c24a..fabdf029885e2b2ac074dd52d60413a25cff2932 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_6368409506294683105+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_6368409506294683105+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_6375780938374445148+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_6375780938374445148+e30acd3a/model.neff
index 5760914f396772e6b8d7877e155f040795d62432..b0c327e000201c9250f089a6209cb81496e63585 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_6375780938374445148+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_6375780938374445148+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_6417297157123190451+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_6417297157123190451+e30acd3a/model.neff
index 863580df2b7244bd40e1f97fdedad0e9a73ab4a1..3ba480067da0fb2f7051265f132bdb52bce87c67 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_6417297157123190451+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_6417297157123190451+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_665145310648895012+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_665145310648895012+e30acd3a/model.neff
index e7a56fa2eb1550c55820bcdc6a304885eae81f9a..6a513cb881f33e8323ff20b3d68a1cf23b8b4a81 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_665145310648895012+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_665145310648895012+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_6711765444274762085+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_6711765444274762085+e30acd3a/model.neff
index 59b570846fb300e1d45a10ee39208e654816c161..98475a02743a622a22e4338e797e9b1385976cac 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_6711765444274762085+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_6711765444274762085+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_7049343467335223052+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_7049343467335223052+e30acd3a/model.neff
index 4a7cb624ee054502e5c600462f230be951602a1c..c2aba7d05d5d5cfccb7221b6b3fd5424f2848025 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_7049343467335223052+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_7049343467335223052+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_7877866862268499659+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_7877866862268499659+e30acd3a/model.neff
index d44af4905d944df114634f951cd21f4b5e2b8930..ccf8bf44a61f6655ceea9e5ba46660a4a0429e58 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_7877866862268499659+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_7877866862268499659+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_793722132614349680+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_793722132614349680+e30acd3a/model.neff
index 54547cff955ddb6481d42465cd4551c8d24192f8..34e57155343c00bf2ecd6af178f14edb2169fd5a 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_793722132614349680+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_793722132614349680+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8142158350842512240+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_8142158350842512240+e30acd3a/model.neff
index 017ebe68b9ef47af003158be1ef24f852d5c8f11..e726b5dbbc99f03a4d0eec71b823c2b46bf904f2 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_8142158350842512240+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_8142158350842512240+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8146958338478347620+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_8146958338478347620+e30acd3a/model.neff
index 944fa0cedd8a992ba39b242b7e8e345e89c133f9..7a3175da41ba05106694e91d223e897f5056f6bd 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_8146958338478347620+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_8146958338478347620+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8245165830758578911+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_8245165830758578911+e30acd3a/model.neff
index a8c0bb21548928a57bb30344f52648781517b86a..e3b019701d678ee60c2ebb9da465e7a4f65a663d 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_8245165830758578911+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_8245165830758578911+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..93276f1d0321eda9f4efbeadfd6a3ed248da3d4f
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_d777179a-9dd6-4706-b29f-9bde032135c4/compiler_workdir/DecoderLayerWrapper/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..be5670c6258a72b3253c262c04b0a3c084a1810d
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3dc99fd8fa13fb76a08ae04c8275958cd87878cfbf86f63ae613d7efcf775bfc
+size 29412
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..3e95b2aa6e481230ebc23f9ec06ccdfe84721b53
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/model.neff
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa8313bb56ae832a4a490d3bb404281b1467e53a380a9ff639c116bbbcb9ede5
+size 328704
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/wrapped_neff.hlo
new file mode 100644
index 0000000000000000000000000000000000000000..011ae30121e455650ad0cc0f22c6191984da0c16
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_831a196add299bcee67a+9235ce67/wrapped_neff.hlo
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:73e6c8d6558ff84236d3ec77d9118741f262bd0c6cc69db8cb09634d171a87c1
+size 334452
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8679096817551502409+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_8679096817551502409+e30acd3a/model.neff
index 18ab2e94ce5ee1cbeadbce69e9a6eb571247596a..e433006c8fedf5ba951489899dee887b682d657c 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_8679096817551502409+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_8679096817551502409+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8882592820567938515+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_8882592820567938515+e30acd3a/model.neff
index 8335d829e64d9ae4699f19575a212a2d0bec5a73..29bda3ac250a4fbdee662d44438fddb644fe9a8c 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_8882592820567938515+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_8882592820567938515+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.hlo_module.pb
index be952da347de317d364c85e97ed45978ec10f0c5..a6adc8e21ca678546dd3249064e046ddd75eb4a5 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.hlo_module.pb
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.hlo_module.pb
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:395529a99f71d2c53ebdea1ddab9d3c20cb88f4a53eef1dc2f557c1d2f43ea83
+oid sha256:d845214d2a88cba3b82d220ba69115b804dd26fbbffef6cf554e026f84ea4e37
 size 474402
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.neff
index 1bc0c8f721ebc3596b98943c0f4a8b48d3e07a43..f23d7a5db2427c619f67fb256c13a16e2e3a95d7 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.neff
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_8d5c7639490b29a58f59+24129607/model.neff
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:be6acce832d0c8333405c5c5a876f2b21d5b6e0729fc0fd863ae38c5fede4fe3
+oid sha256:8f60f1d06a29b43ece2785d79f4c41eb1926fd28863efdc8446127021dd7dc02
 size 42322944
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_9168843780052024308+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_9168843780052024308+e30acd3a/model.neff
index bd9ceabafafd5263517ea7c557613154fae495e0..3c1ef3afb25016fb464f7031c4c568c59878d6aa 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_9168843780052024308+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_9168843780052024308+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_9527558113976496538+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_9527558113976496538+e30acd3a/model.neff
index d53a971069733e5f54b0590696464cdc0680010c..19ba286334ec512636306aee76a5e012bc607871 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_9527558113976496538+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_9527558113976496538+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_9533302140457976101+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_9533302140457976101+e30acd3a/model.neff
index e975fc18a182ded63a0d1883a8a4bd0e813d8406..b4f143f1024fa6a72cb42b979010c607d5b0694a 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_9533302140457976101+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_9533302140457976101+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_9770544877851564228+e30acd3a/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_9770544877851564228+e30acd3a/model.neff
index 1396cead634baf574e7d1075380cf5d1b3c50680..3b6f11ce46f4cbf26c57ca5e0912c05db68d40ab 100644
Binary files a/neuronxcc-2.21.33363.0+82129205/MODULE_9770544877851564228+e30acd3a/model.neff and b/neuronxcc-2.21.33363.0+82129205/MODULE_9770544877851564228+e30acd3a/model.neff differ
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..2d6a49de856d6f86b37d56d7e385723e5c87844c
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_a353f827-4fa1-4e23-b1c0-2eb0986e4f4b/compiler_workdir/DecoderLayerWrapper/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..07bf62460bc7813e4a95f208598e2fa4ae551d41
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76dc525043c8d0fc6b46382b17659e6a2f2e981788c635021a8a61aa832d2866
+size 14480
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..61a6a9150d10ef5d8f83d7638893fb852bd36acb
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/model.neff
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32d578925faca5388c643e94b3afad13b47f62b4c2ca7bbc228bb335c64d19c5
+size 267264
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/wrapped_neff.hlo
new file mode 100644
index 0000000000000000000000000000000000000000..820a9642f67ff9ff8dfa972acb39b2ce7effd4cc
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_ddc85f3b3100f981ecaa+6f338dce/wrapped_neff.hlo
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09eb88c56b73be4ebf969b21a2d19a0143ed5b5a18fa40dc52693c2776269a67
+size 272962
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..b7438c095c70bae643607bcf113d6b18fb357b41
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_a176a9f3-1a1e-438f-a0bd-97d94031e3ee/compiler_workdir/NeuronRMSNorm/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..6a2752c25c9f7e1092e2b55c010694e168610ce3
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c774c6626fc91684e924b843c1fe807dadbbac6d8a8e19443663b43c13eac35
+size 1165
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..78c09adf372bc2031f2e960d76653a58367d5085
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/model.neff
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9926400ecd4630bbc72d37db12f261686af50c967b670e98d324301c618c4a9a
+size 103424
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/wrapped_neff.hlo
new file mode 100644
index 0000000000000000000000000000000000000000..8e8a72b7a93c228e2ebf4a361b5e3eda0a008747
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_e9a55299fdbec8d03417+769ba117/wrapped_neff.hlo
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44b13f663a7a4ea3b224cbd567ad964042c868e4e1e92ea9ea26794eb66d3e20
+size 104320
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.hlo_module.pb
index db90fa54202e67f9c2b63a9f960524c316db768c..087adad3be38716000a37f13a671422143669eca 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.hlo_module.pb
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.hlo_module.pb
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8ed2e09754ffce7b5957680d04d07b9b18166a7809582694d1e73f3af4b21434
+oid sha256:18a82b3beb86fc5ad8db17a90281a3d8d1e499f2379ca89b897d6d0517a715c4
 size 590808
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.neff
index 5bb9136599a471ac2df41827d31feec855e0b172..41c7f9121ab84c7b8d5f62d52bdc7b9ff17da760 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.neff
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/model.neff
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7ec24b7df3cfb5edbbcca57368c392c0044680d9bc3b920a7d250cea571252e6
+oid sha256:ae62b7cc55200c9ef4113fd6a439ddd37d430994cf38d0bdace8cd7f1263e9a4
 size 1547264
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/wrapped_neff.hlo
index 03612fb87e58ac1729b681f05dff4e66f21fcad7..15da5727bb8fa63753b2be337e22a3cb39927dbd 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/wrapped_neff.hlo
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f19ecf2e2bcc3a4b1bd9+a02c3a36/wrapped_neff.hlo
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7abfdf00ed67ea7cd63f8799d9f520f6ca160e13f2f06bcdfebbbc36e860c979
+oid sha256:12dcb6a1b8c66716df9d7b0f72b1f31040c77b92eaede0a8194c2381f5368468
 size 1703455
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..6131411f2bd18de312cdcf224f3a13ec2c95c12c
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_c59a9683-3452-4b5d-9889-165a4d8255d3/compiler_workdir/NeuronLlamaMLP/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..e6084ce6d845ff91b01c0cbfd758d24019f38ef9
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a99355de7043422bc03a8df795afecf586999ec744407bc4fbffe63f002ba2df
+size 1931
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..cf4fb49d20eb48b099a0c724688faab67ca8b324
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/model.neff
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d8afab227390099e5062d23c6d3d05d54d58cc43e081fd3f9809d1b7a32a0c8
+size 134144
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/wrapped_neff.hlo
new file mode 100644
index 0000000000000000000000000000000000000000..575301aa102e817aaebd8602c1dffad5a2788cc3
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f351816884cd2bc34581+f5d1f2e3/wrapped_neff.hlo
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:651f3422f6e0e417c4bc52e7eb2450c292f18513a5651dd9cf53a2b74cfd1efd
+size 136222
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/compile_flags.json b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/compile_flags.json
new file mode 100644
index 0000000000000000000000000000000000000000..08d56ef6ace27afe68cf20bdfb215f041797d6d2
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/compile_flags.json
@@ -0,0 +1 @@
+["--target=trn1", "--enable-saturate-infinity", "--auto-cast=none", "--model-type=transformer", "-O1", "--logfile=/tmp/nxdi_test_ca573acb-72b9-487f-8c42-bb8c33881b46/compiler_workdir/AttentionModelWrapper/_tp0_bk0/log-neuron-cc.txt", "--enable-internal-neff-wrapper"]
\ No newline at end of file
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.done b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.done
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.hlo_module.pb
new file mode 100644
index 0000000000000000000000000000000000000000..8496781848c02ecd308b5b2c5a3c342b702425e7
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.hlo_module.pb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51e79a5cc325f9a317d158d931d75ec73ae1cfb16a4d83343a2593edd95d0044
+size 8979
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.neff
new file mode 100644
index 0000000000000000000000000000000000000000..e6db7e4c58cf454168b4574185e84d5db071a5ec
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/model.neff
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:def7f7f1346fcee0512a90675ebb1ec8330c320dbf15cddc63398b4457c9e6e6
+size 246784
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/wrapped_neff.hlo
new file mode 100644
index 0000000000000000000000000000000000000000..d5a06e4812a4f04b5f95af123da0c5d8afbfc5c4
--- /dev/null
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_f3d762785be23f044440+44e7e4cd/wrapped_neff.hlo
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfacc49a921afcb2e04e2728054de8da9eec3ba49e5eb850328f5eae724f008d
+size 249608
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.hlo_module.pb b/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.hlo_module.pb
index 1f4ba21c9f53777d279e086808276ab93c3b0615..47dc969631220fdb40e5fe10f7bfae7a5454d613 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.hlo_module.pb
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.hlo_module.pb
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6661ed9cdcff742b06a6c5ab009981de926925468498a831dec14ba213178f65
+oid sha256:018a62e9a6d7b944baff314113a3d60b4bee472c3df99216573823afea25c4f2
 size 426072
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.neff b/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.neff
index 969dcbd490fdff3e8bcda5837e8bacb161b67bb1..4676bd1206f5d26f4812a3589bb9f4b6b12cda86 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.neff
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/model.neff
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6047cec9dad7f5b04b2a4ff49b6a2a74235d454de23e4f91a7f6a47fb49a3596
+oid sha256:6948cf6ad7af326ad9824daa8fe53b5c23796a13c29fc690b3bbad9ddae7bef6
 size 2100224
diff --git a/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/wrapped_neff.hlo b/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/wrapped_neff.hlo
index 32c83401835f12a5e8f6eb5ee4e37a3b3ce67a16..937a20a0d55a55d6e55cd233e61afdaa2e032e4d 100644
--- a/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/wrapped_neff.hlo
+++ b/neuronxcc-2.21.33363.0+82129205/MODULE_fde2b1111ad10bd7504e+a02c3a36/wrapped_neff.hlo
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:28f23d333b4867fd00cba757905074f38a5eead20b295e60f78becaf8de1dbf3
+oid sha256:509debbf3be2171f7256bdbf6f813d69a574023b48243f2cf86f759d155740b1
 size 2174311