@@ -50,7 +50,6 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
5050 for (int i = 0 ; i < config .numberOfLayers (); i ++) {
5151 // === Attention Block ===
5252 tornadoForwardScheduler .addWorkerGrid ("layer_" + i + ".attn_rms_reduce" , rmsNormWorker );
53- //tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".attn_rms_apply_fp16", rmsNormWorker);
5453 tornadoForwardScheduler .addWorkerGrid ("layer_" + i + ".qkv_projection" , fusedQKVWorker );
5554 tornadoForwardScheduler .addWorkerGrid ("layer_" + i + ".rope_and_kv_cache" , ropeWithCacheWorker );
5655 tornadoForwardScheduler .addWorkerGrid ("layer_" + i + ".attention" , parallelAttentionWorker );
@@ -202,21 +201,6 @@ TaskGraph setupSingleFFNLayer(LlamaTornadoWeights weights, Configuration config,
202201 TransformerComputeKernelsLayered ::reductionOneBlockWithLayerFuseFP16 ,
203202 context , state .wrapXbFP16 , state .wrapX , weights .rms_att_weightLayered [layerIndex ].asFloatArray (), state .temp ,
204203 config .dim (), config .rmsNormEps (), state .localSize );
205- /*unifiedLayer.task("attn_rms_reduce",
206- TransformerComputeKernelsLayered::reductionOneBlockWithLayer,
207- context, state.temp, state.wrapX,
208- config.dim(), config.rmsNormEps(), state.localSize);
209-
210- if (shouldUseFinalNormalization()) {
211- unifiedLayer.task("attn_rms_finalize",
212- TransformerComputeKernelsLayered::reductionFinalNormalization,
213- context, state.temp, config.dim(), config.rmsNormEps());
214- }
215-
216- unifiedLayer.task("attn_rms_apply_fp16",
217- TransformerComputeKernels::mapContextWithQuantize,
218- context, state.wrapXbFP16, state.wrapX,
219- weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.temp);*/
220204
221205 // QKV Projection (fused)
222206 unifiedLayer .task ("qkv_projection" ,
0 commit comments