User problem
Logit materialisation is immensely expensive for long context. Multiple good LCE implementations exist to support this (Cross Cut Entropy, Liger, etc.). They also appear to support softmax capping out of the box for models like Gemma4. These are much more memory efficient than the manual method, and I expect should be faster as more work is done in the same kernel (although I haven't confirmed this).
These are supported by Automodel so this is a case where Automodel outperforms Bridge on memory usage.
I note there has been discussion about a fused CE being implemented in Megatron/TE for the last 6+ months but so far has only landed in dev, and has not been as widely tested.
Desired outcome
Add support for a fused CE implementation and a flag to enable it in the more common models, perhaps via a uv extra.
Alternatives considered
No response
Affected area
area:model
Urgency / use case
Important but not blocking
Extra context
No response
User problem
Logit materialisation is immensely expensive for long context. Multiple good LCE implementations exist to support this (Cross Cut Entropy, Liger, etc.). They also appear to support softmax capping out of the box for models like Gemma4. These are much more memory efficient than the manual method, and I expect should be faster as more work is done in the same kernel (although I haven't confirmed this).
These are supported by Automodel so this is a case where Automodel outperforms Bridge on memory usage.
I note there has been discussion about a fused CE being implemented in Megatron/TE for the last 6+ months but so far has only landed in dev, and has not been as widely tested.
Desired outcome
Add support for a fused CE implementation and a flag to enable it in the more common models, perhaps via a uv extra.
Alternatives considered
No response
Affected area
area:model
Urgency / use case
Important but not blocking
Extra context
No response