Feature Request: Enable Initial Base Support for MLP Speculators

[markurtz]

Summary

Implement and finalize src/speculators/models/mlp.py to support MLP-based speculator models for speculative decoding algorithms utilizing combined token/embedding speculators.

Reference

• Accelerating Production LLMs with Combined Token/Embedding Speculators

Acceptance Criteria

Classes and Test Cases

• Implement MLPSpeculatorConfig and MLPSpeculator following the example in src/speculators/models/eagle.py.
• Ensure compatibility with SpeculatorModelConfig.from_pretrained and SpeculatorModel.from_pretrained.
• Implement comprehensive test cases following the examples in:
  • tests/unit/models/test_eagle_config.py
  • tests/unit/models/test_eagle_model.py

MLPSpeculatorConfig

• Include all relevant hyperparameters expected to change or be configured to construct a working MLP Speculator model as defined in the referenced paper.

MLPSpeculator

• Correctly create the required MLP-based architecture from a given MLPSpeculatorConfig.
• Enable loading and saving of weights.
• Integrate seamlessly with the existing system.

TokenProposal Functionality

• Implement any missing TokenProposal methods or functionality as defined in the paper, or expand current implementations as needed.

Out of Scope (Future Targets)

• Implement a functioning forward pass for SpeculatorModel compatible with training flows.
• Implement a functioning generate pass for SpeculatorModel compatible with generation flows.
• Create an Algorithm factory to handle preconfigured hyperparameters for the desired supported algorithms.

[thisisatharva-rh]

Happy to take this up

[yuchenwang3]

I’d like to work on this—could you assign me?

[thisisatharva-rh]

Hey Wangyu, sure! I was working on this prior to drowning with other stuff. let me know if you need any help!