Skip to content

[sglang-miles] Support INT4 QAT for RL - #18565

Merged
yueming-yuan merged 1 commit into
sgl-project:sglang-milesfrom
yueming-yuan:sglang-miles
Feb 12, 2026
Merged

[sglang-miles] Support INT4 QAT for RL#18565
yueming-yuan merged 1 commit into
sgl-project:sglang-milesfrom
yueming-yuan:sglang-miles

Conversation

@yueming-yuan

@yueming-yuan yueming-yuan commented Feb 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Summary

  • Add asymmetric (non-symmetric) INT4 quantization support for MoE models with zero-point parameters
  • Add post_process_weights API for Marlin↔GPTQ weight format conversion during RL training rollout
  • Add zero-point repacking for Marlin format via moe_awq_to_marlin_zero_points
  • Skip weight transposition for zero-point parameters in MoE weight loader

Changes

  • compressed_tensors_moe.py: Support asymmetric quantization with zero-point creation, Marlin zero-point repacking, restore_weights_before_loading for rollback
  • compressed_tensors.py: Remove is_symmetric requirement from _is_wNa16_group_channel
  • fused_moe_triton/layer.py: Skip transposition for zero-point weights
  • io_struct.py: Add PostProcessWeightsReqInput/ReqOutput dataclasses
  • engine.py, http_server.py, scheduler.py, scheduler_update_weights_mixin.py, tokenizer_communicator_mixin.py, tp_worker.py, model_runner.py: Wire up post_process_weights API through the full request pipeline

@gemini-code-assist

gemini-code-assist Bot commented Feb 10, 2026

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@yueming-yuan yueming-yuan changed the title Support non-symmetric INT4 QAT for MoE models [sglang-miles] Support INT4 QAT for RL Feb 10, 2026
@yueming-yuan
yueming-yuan merged commit 5ec6118 into sgl-project:sglang-miles Feb 12, 2026
1 check passed
yueming-yuan added a commit that referenced this pull request Feb 13, 2026
@yueming-yuan
[sglang-miles] Support INT4 QAT for RL (#18565)
7adbd21
yueming-yuan added a commit that referenced this pull request Mar 5, 2026
@yueming-yuan
[3/8] [sglang-miles] Support INT4 QAT for RL (#18565)
19e8de4
yueming-yuan added a commit that referenced this pull request Mar 27, 2026
@yueming-yuan
[3/8] [sglang-miles] Support INT4 QAT for RL (#18565)
1e764da
yueming-yuan added a commit that referenced this pull request Apr 6, 2026
@yueming-yuan
[3/8] [sglang-miles] Support INT4 QAT for RL (#18565)
2bc3f68
zyzshishui pushed a commit to zyzshishui/sglang that referenced this pull request Apr 25, 2026
@yueming-yuan @zyzshishui
[3/8] [sglang-miles] Support INT4 QAT for RL (sgl-project#18565)
6e57cdf
@zyzshishui zyzshishui mentioned this pull request Apr 25, 2026
Closed
5 tasks
yueming-yuan added a commit that referenced this pull request May 21, 2026
@yueming-yuan
[3/14] [sglang-miles] Support INT4 QAT for RL (#18565)
5bd6c6b
This was referenced Jun 17, 2026
Closed
Merged
JessicaJiang-123 pushed a commit to JessicaJiang-123/sglang that referenced this pull request Jun 21, 2026
@yueming-yuan
[3/14] [sglang-miles] Support INT4 QAT for RL (sgl-project#18565)
b2cfcc1
nanjiangwill pushed a commit to nanjiangwill/sglang that referenced this pull request Jul 7, 2026
@yueming-yuan @JD-ETH @maocheng23
[sglang-miles] RL weight processing: begin/end weight-update sessions (…
4c7ee89 
…sgl-project#18565, sgl-project#22663, sgl-project#28001, sgl-project#29675)

Combined weight-processing machinery, folding the general post-process
plumbing from the INT4 QAT work (sgl-project#18565), the P2P post_load_weights
extension (sgl-project#22663), the weight-processing refactor (sgl-project#28001), and the
pause-aware weight locking fix (sgl-project#29675) into its final form:

- Begin/EndWeightUpdate request chain: engine/http_server/
  tokenizer_control_mixin (pause-aware locking) -> scheduler ->
  weight_updater sessions with {target,draft,all} runner selectors
- loader.py: public post_load_weights, restore_weight/postprocess_weight
  via _apply_quant_method_hook (skips LoRA wrappers)
- model_runner begin/end_weight_update + receive_weights_from_distributed
- iter_runners() on tp_worker and all spec v2 workers (incl. port to
  DFlashWorkerV2; v1 dflash/frozen-kv workers no longer exist)
- CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch
- weight_checker ComparableWeight refactor + weight_checker_comparator

Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 14, 2026
@yueming-yuan @JD-ETH
@maocheng23 @guapisolo
[18/24] [sglang-miles] RL weight processing: begin/end weight-update …
e934d92 
…sessions + spec-draft weight check (#18565, #22663, #28001, #29675, #27750)

Combined weight-processing machinery, folding the general post-process
plumbing from the INT4 QAT work (#18565), the P2P post_load_weights
extension (#22663), the weight-processing refactor (#28001), and the
pause-aware weight locking fix (#29675) into its final form, plus the
extension of the weight checker to speculative draft worker(s) (#27750):

- Begin/EndWeightUpdate request chain: engine/http_server/
  tokenizer_control_mixin (pause-aware locking) -> scheduler ->
  weight_updater sessions with {target,draft,all} runner selectors
- loader.py: public post_load_weights, restore_weight/postprocess_weight
  via _apply_quant_method_hook (skips LoRA wrappers)
- model_runner begin/end_weight_update + receive_weights_from_distributed
- iter_runners() on tp_worker and all spec v2 workers (incl. DFlashWorkerV2)
- CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch
- weight checker: selector + skip_tensor_list + role-prefixed overall_checksum
  over get_model_runners(selector), replacing _get_draft_model_runner

Kept v0.5.15's upstreamed ComparableWeight weight_checker_comparator
(incl. the #29623 ue8m0 scale-unpack fix); adapted the weight_checker
caller to its public compare_weights/CHUNK_NUMEL API.

Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Co-authored-by: Jiajun Li <jiajun.li@radixark.ai>
yueming-yuan added a commit that referenced this pull request Jul 14, 2026
@yueming-yuan @JD-ETH
@maocheng23 @guapisolo
[18/24] [sglang-miles] RL weight processing: begin/end weight-update …
130048f 
…sessions + spec-draft weight check (#18565, #22663, #28001, #29675, #27750)

Combined weight-processing machinery, folding the general post-process
plumbing from the INT4 QAT work (#18565), the P2P post_load_weights
extension (#22663), the weight-processing refactor (#28001), and the
pause-aware weight locking fix (#29675) into its final form, plus the
extension of the weight checker to speculative draft worker(s) (#27750):

- Begin/EndWeightUpdate request chain: engine/http_server/
  tokenizer_control_mixin (pause-aware locking) -> scheduler ->
  weight_updater sessions with {target,draft,all} runner selectors
- loader.py: public post_load_weights, restore_weight/postprocess_weight
  via _apply_quant_method_hook (skips LoRA wrappers)
- model_runner begin/end_weight_update + receive_weights_from_distributed
- iter_runners() on tp_worker and all spec v2 workers (incl. DFlashWorkerV2)
- CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch
- weight checker: selector + skip_tensor_list + role-prefixed overall_checksum
  over get_model_runners(selector), replacing _get_draft_model_runner

Kept v0.5.15's upstreamed ComparableWeight weight_checker_comparator
(incl. the #29623 ue8m0 scale-unpack fix); adapted the weight_checker
caller to its public compare_weights/CHUNK_NUMEL API.

Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Co-authored-by: Jiajun Li <jiajun.li@radixark.ai>
yueming-yuan added a commit that referenced this pull request Jul 14, 2026
@yueming-yuan @JD-ETH
@maocheng23 @guapisolo
[18/24] [sglang-miles] RL weight processing: begin/end weight-update …
3ad9ca3 
…sessions + spec-draft weight check (#18565, #22663, #28001, #29675, #27750)

Combined weight-processing machinery, folding the general post-process
plumbing from the INT4 QAT work (#18565), the P2P post_load_weights
extension (#22663), the weight-processing refactor (#28001), and the
pause-aware weight locking fix (#29675) into its final form, plus the
extension of the weight checker to speculative draft worker(s) (#27750):

- Begin/EndWeightUpdate request chain: engine/http_server/
  tokenizer_control_mixin (pause-aware locking) -> scheduler ->
  weight_updater sessions with {target,draft,all} runner selectors
- loader.py: public post_load_weights, restore_weight/postprocess_weight
  via _apply_quant_method_hook (skips LoRA wrappers)
- model_runner begin/end_weight_update + receive_weights_from_distributed
- iter_runners() on tp_worker and all spec v2 workers (incl. DFlashWorkerV2)
- CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch
- weight checker: selector + skip_tensor_list + role-prefixed overall_checksum
  over get_model_runners(selector), replacing _get_draft_model_runner

Kept v0.5.15's upstreamed ComparableWeight weight_checker_comparator
(incl. the #29623 ue8m0 scale-unpack fix); adapted the weight_checker
caller to its public compare_weights/CHUNK_NUMEL API.

Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Co-authored-by: Jiajun Li <jiajun.li@radixark.ai>
yueming-yuan added a commit that referenced this pull request Jul 14, 2026
@yueming-yuan @JD-ETH
@maocheng23 @guapisolo
[18/24] [sglang-miles] RL weight processing: begin/end weight-update …
6f950d9 
…sessions + spec-draft weight check (#18565, #22663, #28001, #29675, #27750)

Combined weight-processing machinery, folding the general post-process
plumbing from the INT4 QAT work (#18565), the P2P post_load_weights
extension (#22663), the weight-processing refactor (#28001), and the
pause-aware weight locking fix (#29675) into its final form, plus the
extension of the weight checker to speculative draft worker(s) (#27750):

- Begin/EndWeightUpdate request chain: engine/http_server/
  tokenizer_control_mixin (pause-aware locking) -> scheduler ->
  weight_updater sessions with {target,draft,all} runner selectors
- loader.py: public post_load_weights, restore_weight/postprocess_weight
  via _apply_quant_method_hook (skips LoRA wrappers)
- model_runner begin/end_weight_update + receive_weights_from_distributed
- iter_runners() on tp_worker and all spec v2 workers (incl. DFlashWorkerV2)
- CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch
- weight checker: selector + skip_tensor_list + role-prefixed overall_checksum
  over get_model_runners(selector), replacing _get_draft_model_runner

Kept v0.5.15's upstreamed ComparableWeight weight_checker_comparator
(incl. the #29623 ue8m0 scale-unpack fix); adapted the weight_checker
caller to its public compare_weights/CHUNK_NUMEL API.

Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Co-authored-by: Jiajun Li <jiajun.li@radixark.ai>
@yushengsu-thu yushengsu-thu mentioned this pull request Jul 17, 2026
Open
5 tasks
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/23] [sglang-miles] RL weight-update sessions + distributed update…
ca22702 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/24] [sglang-miles] RL weight-update sessions + distributed update…
1f9d84a 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/25] [sglang-miles] RL weight-update sessions + distributed update…
57ef29a 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/27] [sglang-miles] RL weight-update sessions + distributed update…
bb17e99 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@Fridge003 Fridge003 Awaiting requested review from Fridge003 Fridge003 is a code owner
@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner
@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner
@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner
@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner
@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1 Edwardf0t1 is a code owner
@FlamingoPg FlamingoPg Awaiting requested review from FlamingoPg FlamingoPg is a code owner
@AniZpZ AniZpZ Awaiting requested review from AniZpZ AniZpZ is a code owner
@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner
@CatherineSue CatherineSue Awaiting requested review from CatherineSue CatherineSue is a code owner
@slin1237 slin1237 Awaiting requested review from slin1237 slin1237 is a code owner
@JustinTong0323 JustinTong0323 Awaiting requested review from JustinTong0323 JustinTong0323 is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@yueming-yuan