[Feature] Make submodules’ dtype be configurable · sgl-project/sglang-jax#661

2026-01-15T04:52:03.000Z

### **Checklist** - [ ] If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed. - [ ] Please use English, otherwise it will be closed. ### **Motivation** Control the precision scope when using higher precision, like float32. * softmax * proj in attn and mlp ### **Related resources**

(11 comments) (0 reactions) (1 assignee)Python (276 stars) (101 forks)auto 404

Featurehelp wanted

Description

Checklist

If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
Please use English, otherwise it will be closed.

Motivation

Control the precision scope when using higher precision, like float32.

softmax
proj in attn and mlp

Related resources

Contributor guide

Tech stack: python
Domain: backend
Issue type: feature
Difficulty: 3
Estimated time: half day
Activity status: active
Clarity: clear
Prerequisites: PythonJAX
Newbie friendliness: 60
Research direction: Investigate how to make the dtype configurable for submodules like softmax, attention projection, and MLP projection to support higher precision (e.g., float32) while maintaining performance.