[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 commenti) (0 reazioni) (1 assegnatario)Python (276 star) (101 fork)auto 404

Featurehelp wanted

Descrizione

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

Guida contributor

Tech stack: python
Dominio: backend
Tipo issue: feature
Difficoltà: 3
Tempo stimato: half day
Stato attività: active
Chiarezza: clear
Prerequisiti: PythonJAX
Adatta ai principianti: 60
Direzione di ricerca: Indaga come rendere configurabile il dtype per i sottomoduli come softmax, proiezione dell'attenzione e proiezione MLP per supportare una precisione maggiore (es. float32) mantenendo le prestazioni.