Refactor: Improve model classes and remove redundant parameters in gemma/model.py

[CodersAcademy006]

This issue tracks two code quality improvements for gemma/model.py to enhance maintainability and API clarity.

1. Redundant Code in Linear and Embedding Classes

• Problem: The Linear and Embedding classes contain nearly identical boilerplate code for handling quantized weights. This violates the DRY (Don't Repeat Yourself) principle.
• Proposed Solution: Refactor this shared logic into a new QuantizedWeight base class. This will make the code cleaner, more modular, and easier to maintain.

2. Unused kv_write_indices Parameter

• Problem: The GemmaForCausalLM.forward method accepts a kv_write_indices parameter that is immediately overwritten by input_positions. The passed argument is never used.
• Proposed Solution: Remove this redundant parameter from the method signature to clean up the API and avoid confusion for developers.

These changes will improve the overall code quality and maintainability without altering the model's functionality.

[Balakrishna-Chennamsetti]

Hi @CodersAcademy006 ,

Thank you for the notification and for your generous contribution.
Thank you so much for your contribution and for reporting the issue. We appreciate you submitting this PR. Your patience while it is reviewed is greatly valued.

Thanks.

[rehan243]

Quick q: are you guys running this with any custom quantization backends, or just sticking to the default torch bits? Noticed the same Linear/Embedding boilerplate when I was poking around gemma/model.py last week — honestly, we had to hack something similar for a LoRA experiment and it got ugly fast. Pulling the quantized weight logic into a shared QuantizedWeight base totally makes sense, especially if you ever end up introducing new layers with quantized params later (been there, it’s a pain to refactor after the fact).

On the kv_write_indices thing: yeah, we hit this exact param weirdness in another repo, kept confusing devs who expected it to actually do something. Cleaning up the forward signature is definitely the way to go, especially for folks digging through the API. Outta curiosity, any reason it stuck around — maybe some legacy checkpoint support or just drifted in over time?

Are you planning to touch any of the config files too, or just model.py for now?