[Bugfix] Fix Per-Token Dynamic Activation Quantization

[max410011]

Summary

This PR fixes the activation quantization issue described in Issue #394, where the input scale shape was incorrect when using the Dynamic TOKEN strategy.

Fix

• Corrected the reduction dimensions to ensure only the hidden dimension is reduced.
• This ensures the input scale shape is (batch_size, seq_len, 1) instead of (1, seq_len, hidden_dim).

[brian-dellabetta]

Hi @max410011 , appreciate the thorough detail in the issue! I tried your PR, and both original main and your branch seem to work, the resultant models can be loaded up and run in vllm, which surprises me. This is some old code, and per-token/per-channel always slips me up. I will ask around to see if your reasoning in the issue description is correct.

[brian-dellabetta]

i validated that this gives the shape described in #394 , and after internal conversations this is correct. This is only an issue when running outside of vllm

[kylesayrs]

Could you add a test to demonstrate and verify that these changes are correct? Awesome catch and resolution, thanks!

[dsikka]

@kylesayrs