Skip to content

[diffusion] avoid redundant LoRA scaling while preserving exact numerics - #43028

Merged
mickqian merged 6 commits into
sgl-project:mainfrom
niehen6174:diffusion-lora-unit-scale
Oct 13, 2026
Merged

[diffusion] avoid redundant LoRA scaling while preserving exact numerics#43028
mickqian merged 6 commits into
sgl-project:mainfrom
niehen6174:diffusion-lora-unit-scale

Conversation

@niehen6174

@niehen6174 niehen6174 commented Oct 8, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Collaborator

Avoid redundant dynamic LoRA scale multiplications while preserving the current quality-tier contract. Unit scales are skipped in every tier; exact and requests without a forward context retain the original order of nonunit strength and request-scale multiplications. Combining those factors is restricted to requests that allow lossless.

The shared scaling helper is integrated with main’s dynamic-LoRA delta helper. Existing fused GEMM updates and row-parallel bias handling remain intact. The alpha/rank factor stays separate in the fallback path.

Regression coverage checks scale-operation counts by quality tier and compares exact outputs bit for bit on CPU and CUDA, including nonbinary strength/request scales (0.7 and 0.6). Targeted remote tests: 77 passed; scoped pre-commit passed. Final CI is tracked on the new head. No GT or accuracy thresholds were changed.

CI States

Latest PR Test (Base): ✅ Run #38173797942
Latest PR Test (Extra): ❌ Run #38173797703
Latest PR Test (AMD ROCm 10): ❌ Run #38173797874

@niehen6174
[diffusion] fix: skip the LoRA delta multiply at unit runtime scale
051cf1c 
sgl-project#30487 folded the per-request LoRA scale into every LoRA linear as a second
tensor multiply. At the default scale of 1.0 it only adds a kernel per call;
LTX-2/2.3 refinement steps run ~1.4-1.7k LoRA linears each and get ~6% slower.
Outputs are bit-identical at every scale.
@github-actions github-actions Bot added lora diffusion SGLang Diffusion labels Oct 8, 2026
@niehen6174
[diffusion] fix: fold the LoRA strength and runtime scale into one mu…
6d9f811 
…ltiply

Multiply the delta by strength * runtime_lora_scale once, and skip it when
the product is 1.0. LTX-2/2.3 run their stage-2 distilled LoRA at strength
1.0, so refinement steps launch no scale multiply at all. Outputs are
bit-identical whenever the runtime scale is 1.0.
@niehen6174 niehen6174 changed the title [diffusion] fix: skip the LoRA delta multiply at unit runtime scale [diffusion] fix: multiply the LoRA delta by its scale at most once Oct 8, 2026
@niehen6174
niehen6174 marked this pull request as ready for review October 8, 2026 03:50
@niehen6174 niehen6174 added the run-ci CI: run the baseline test suite on this PR label Oct 8, 2026
detain added a commit to detain/sglang that referenced this pull request Oct 8, 2026
@detain
diffusion: share/fold LoRA delta scaling (sgl-project#43028, ported)
fac65fa 
sgl-project#43028's base carries the old 4-site LoRA forwards plus the runtime scale
input from sgl-project#30487; this stack has the sgl-project#37460 rework and no runtime scale,
so the three commits land as their net final state: BaseLayerWithLoRA.
_scale_lora_delta folds alpha/rank and strength multiplies, skipping
elementwise muls entirely at unit scale (measured ~6%/multiply on LTX-2
refinement LoRA stages). Ported test counts aten.mul ops at 0/1 across
strength x {nn.Linear, ReplicatedLinear}.

PR commits: 051cf1c 6d9f811 17302ee
@mickqian mickqian changed the title [diffusion] fix: multiply the LoRA delta by its scale at most once [diffusion] avoid redundant LoRA scaling while preserving exact numerics Oct 11, 2026
@mickqian
mickqian merged commit feba04f into sgl-project:main Oct 13, 2026
134 of 148 checks passed
detain added a commit to detain/sglang that referenced this pull request Oct 15, 2026
@detain
Merge upstream/main (337 commits) into qwennext-stack — merge sgl-pro…
8639319 
…ject#4

Conflict resolutions (35 files), preserving stack behavior:

srt/utils/common.py: make_layers — folded our per-layer empty_cache body
  INTO upstream's stage-binding build() (layer_stack refactor sgl-project#43453-43460).
  Upstream removed is_building_neighbour_layer/building_neighbour_layer
  entirely (our earlier sgl-project#42481 pick); only user was qwen4_exp.
srt/models/qwen4_exp.py: import unions (cached_property/partial/TYPE_CHECKING,
  linear_attn_parallel_group); dropped dead is_building_neighbour_layer import;
  prefetch_stream keeps our _graph_prefetch_buffer line with upstream's
  ple_offload_embedding conditional; our PLE host-staging design kept
  (_offloads_ple guard moot under new stage facts).
srt/managers/cache_controller.py: upstream sgl-project#43550 absorbed our sgl-project#39830
  prefetch-degrade fixes; took upstream restructure, deleted auto-merge's
  stale duplicate tails (double-ack/deadlock hazard), kept our rationale
  comments. Fenced start_writing + load_fence_stream survive.
srt/mem_cache/hybrid_cache/hybrid_cache_controller.py: same pattern;
  sidecar degrade comment updated ('upstream adopted unconditionally').
model_executor/pool_configurator.py: UNION draft-cell sizing — our FP4
  _compute_eagle_fp4_draft_cell_size branch + upstream HiSparse
  draft_loc_space_scale else-branch. b05a7c1 None layer_info guard intact.
srt/layers/attention/qsa/sparse_attn.py: kept OUR FP8-descale chunk loop
  (superset; upstream sgl-project#41933 only re-commented); took upstream
  _pack_qsa_prefill_kv + pack_qsa_prefill_kv; our ck interface keeps
  k_scale/v_scale/max_q; __all__ union.
qsa/mqa.py: import functools + min_scored_len kwarg into our triton-first
  _resolve_mqa_backend dispatch; upstream ctas_per_row rework kept.
qsa/qsa_indexer.py: batch-invariance gate moved INTO upstream's shared
  use_fast_topk predicate.
qwen_sparse_attn_backend.py: imports union; gather path UNION — upstream
  pack_qsa_prefill_kv used when scales are unit (QSA stores unscaled K/V),
  our index_select fallback when checkpoint k/v descales != 1.0.
managers/schedule_policy.py: add_chunked_req UNION — upstream slot-reserve
  gate first, then our prefill_chunk_alignment truncation (monotone-safe);
  all 3 alignment sites (1177/1340/1608) intact.
managers/tokenizer_manager.py: import unions (PositionalEmbedsTransport +
  RequestPreprocessor + watermarking fns).
entrypoints/openai/serving_base.py: kept OUR request_preprocessor.run path.
model_executor/runner/eager_runner.py: took upstream TCPCG eager-fallback
  removal (sgl-project#42287); kept our shared cp_extend_forward call shape (our
  runner_utils/cp_extend.py is byte-equivalent to upstream's
  _execute_extend_cp — single source of truth).
runner/flashinfer_autotune.py: import union (can_run_dummy_draft_extend +
  cdiv/ceil_align + upstream get_eager_max_batch_size sgl-project#43789).
arg_groups/model_hook.py: Gemma4 accepted_backends union fa4 + flashinfer
  (NVFP4-KV enablement).
multimodal_gen/runtime/layers/lora/linear.py: KEPT OUR sgl-project#37460 pipeline
  structure; upstream runtime_lora_scale==0 early-return ordering adopted;
  _add_lora_delta/_add_lora_output_offset wired into pipeline; kept
  UPSTREAM's quality-tiered _scale_lora_delta (deleted our simple variant).
  DEBT: upstream RowParallel bias_in_gemm lossless opt not ported (perf-only).
multimodal_gen/runtime/layers/attention/layer.py: _gather_sharded_sequence
  took upstream sgl-project#43793 .contiguous() fix (subsumes our 55aaa02 guard);
  our _forward_with_kv_gather k/v guards kept (separate path, not covered).
configs/pipeline_configs/base.py + configs/sample/minimax_h3.py + envs.py:
  unions (supports_continuous_batching / teacache_params+fields /
  SGLANG_TEACACHE_CALIBRATE).
models/dits/ltx_2.py mask-fill: took upstream (identical to our fix).
ltx_2/denoising.py: took upstream pass_specs (sgl-project#43667) — subsumes our
  legacy one-stage branch; our _ltx2_model_forward_context helpers already
  upstreamed (entry-142 sgl-project#42928 CONFIRMED SUPERSEDED).
ltx_2/decoding_av.py: KEPT OUR split-stage AV decode (independent pool
  scaling) over upstream single-class restructure; adopted upstream
  chunked pinned-memory _postprocess_video; TODO(merge-4): upstream
  EncodeWhileDecoding streaming + rank-0 retain gating not ported.
kernels/ops/diffusion/activation/silu_mul_bitexact.py: our out= column-slice
  support + XPU guard UNION with upstream @replay_launches (sgl-project#43859).
kernels/ops/diffusion/README.md + docs performance-optimization.mdx: unions.
rust/sglang-renderer: ACCEPTED upstream crate removal (sampling.rs git rm);
  our sgl-project#42779 top_k vocab-bound port moves to sglang-server message/sampling.rs
  (follow-up).
Tests: anthropic test_serving — UNION fake-class exception param + our
  request_preprocessor stubs, took upstream's _FakeNonStreamingOpenAI removal
  (68 pass). test_server_args — union both new tests (pre-existing GPU-only
  collection ENV fail unchanged). test_lora_inference_mode — took upstream
  sgl-project#43028 runtime_scale+quality-tier test (supersedes our adapted variant;
  our kept quality-tiered _scale_lora_delta satisfies it).
  test_fused_moe_triton_config — kept our int4-filename + RTX-PRO-6000 JSON
  tests despite upstream sgl-project#43503 test-drop. test_hicache_file_lru_unit —
  union our corrupt-page-degrade class + upstream sgl-project#42077 uncached-batch test
  (40 pass). test_hicache_host_register — upstream alias host_common.
  test_cuda_graph_setup — took upstream rewrite of our override-decode-runner
  test into draft-local-capture-sizes test (sgl-project#43303 API).

Verified auto-merged stack markers survived: flashinfer_backend KVLocPlan/
prefill_wrappers_ragged/NVFP4-KV forwarding, memory_pool set_kv_buffer_prefix_valid,
fp8 grouped_moe kwarg, fp8_utils reshape fix 3642490, SGLANG_QWEN4_PLE_FILE_REUSE,
schedule_batch freed_rows fix, prefill_chunk_alignment sites, sgl-project#42736 autotune
default-on, PCIe-IPC vocab gathers, flashinfer 0.7.1 pin taken (NOT installed).

Co-Authored-By: Claude Opus 5.5 (1M context)
detain added a commit to detain/sglang that referenced this pull request Oct 15, 2026
@detain
fix(diffusion): restore lost RowParallelLinearWithLoRA.forward qualit…
3f077d4 
…y-gated bias path

Merge sgl-project#4 kept our stack's pre-sgl-project#43501 class body and dropped upstream's new
forward override (split through the generic _forward_with_delta
decomposition, which always fuses the bias in the GEMM). That silently
changed exact-tier numerics (bias added after reduce, bit-for-bit reference
chain) and broke upstream's own test_row_parallel_bias_moves_into_the_gemm_
only_above_exact (new in sgl-project#43028) on every host, including GPU CI. Restores
the override verbatim (imports split_tensor_along_last_dim and
tensor_model_parallel_all_reduce) while keeping the stack's __init__,
rank-0-gated _add_lora_output_offset, and slice hooks. Test file now matches
origin/main and passes 25P/20S CPU.

(cherry picked from commit a600fb88f46b88cd599f6dd7925891ba1842a28a)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mickqian mickqian Awaiting requested review from mickqian mickqian is a code owner
@ping1jing2 ping1jing2 Awaiting requested review from ping1jing2 ping1jing2 is a code owner
@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner
@yichiche yichiche Awaiting requested review from yichiche yichiche is a code owner
@AgainstEntropy AgainstEntropy Awaiting requested review from AgainstEntropy AgainstEntropy is a code owner
@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner
@kevin-mii kevin-mii Awaiting requested review from kevin-mii kevin-mii is a code owner

Assignees

No one assigned

Labels

diffusion SGLang Diffusion lora run-ci CI: run the baseline test suite on this PR

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@niehen6174 @mickqian