[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!

[kylesayrs]

Shouldn't this be generalized to reflect all dims except the last? There are cases where activations are passed with 4 or 5 dimensions, not just 3

[max410011]

As far as I know, the input/output to a linear layer in typical LLMs usually has the shape (bs, seq_len, hidden_dim).

Other activation shapes, such as (bs, num_heads, seq_len, head_dim), generally appear in the attention computation (e.g., Q @ K^T, attention_weights @ V). If we intend for this function to support quantization in those parts as well, then yes, it makes sense to generalize it accordingly.