minor: read attn impl from config json

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/eagle_config.json

@@ -6,5 +6,6 @@
         "original_max_position_embeddings": 8192,
         "rope_type": "llama3"
     },
-    "initializer_range": 0.02
+    "initializer_range": 0.02,
+    "attn_implementation": "flex_attention"
 }

modelopt/torch/speculative/plugins/transformers.py

@@ -184,10 +184,6 @@ def __init__(self, config, decoder_layer_cls, bias=False):
         super().__init__()
         self.config = config
 
-        # Use flex attention for efficient TTT
-        # config._attn_implementation = "flex_attention"
-        config.attn_implementation = "sdpa"
-
         self.layers = nn.ModuleList(
             [decoder_layer_cls(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
         )
@@ -446,7 +442,6 @@ def modify(
             eagle_architecture_config=eagle_architecture_config,
         )
         self.eagle_config = PretrainedConfig.from_dict(eagle_architecture_config)
-        self.eagle_config._attn_implementation = "sdpa"
         decoder_cls = (
             type(self.model.layers[-1]) if self.eagle_reuse_base_decoder else LlamaDecoderLayer
         )