Conversation
Pack HC weights once with AITER during the loader quant_method post-load phase and replace the original storage. Share the packed layout between B1..32 FlyDSL down/up kernels and a layout-aware torch.compile fallback. Prepare workspaces per module, physical B and stream before capture. Add SGLANG_GR_READ_FLYDSL (default on) to select the integration or the original path at model construction. Preserve norm/residual/combine contracts and reject unsafe packed-weight updates or device/dtype moves. Forward profiler stack/shape environment settings to the profile API.
Comment on lines
+51
to
+64
| def mix_packed(x, packed_down, packed_up): | ||
| """Original Torch math, with weight layout decoding inside the graph. | ||
|
|
||
| Dense intermediates are call-local. In particular, do not cache these | ||
| decoded weights: both prefill and FlyDSL own only the packed storage. | ||
| """ | ||
| w_down = unpack_weight(packed_down, R, K) | ||
| w_up = ( | ||
| unpack_weight(packed_up, K, R).reshape(HS, HC, R).permute(1, 0, 2).reshape(K, R) | ||
| ) | ||
| gates = F.silu(F.linear(x, w_down) / HC) | ||
| gates = F.linear(gates, w_up) | ||
| gates = torch.sigmoid(gates).unflatten(-1, (HC, HS)) | ||
| return (gates * x.unflatten(-1, (HC, HS))).mean(dim=-2) |
Collaborator
There was a problem hiding this comment.
One dummy question, for a server running with BS 1-64 in case prefix cache enabled.
We have flydsl GR read and original torch compile kernel, do we trigger weight preshuffle for BS 1-32 via pack_weights and unpack_weight for BS > 32 during runtime?
| @@ -0,0 +1,93 @@ | |||
| # ROCm Qwen GR read | |||
Collaborator
There was a problem hiding this comment.
Let's remove markdown after PR ready. SGLANG_GR_READ_FLYDSL=1 should be mentioned in PR for enable this opt.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Use flydsl for decode (batch 1-32), fallback to torch.compile for others. Introduce preshuffle weight to support flydsl gemm operations.
Add
SGLANG_GR_READ_FLYDSL(default on) to select the integration or the original path at model construction.Motivation
Modifications
packed_weightsapply preshuffling to both down gemm and up gemm weights viaaiter.ops.shuffle.shuffle_weightFor up gemm weights, perform an extra hc/hidden dimension conversion (up_interleaved[h*4+c,r] = w_up[c*2560+h,r])before the preshuffle step.Accuracy Tests
Speed Tests and Profiling
profiled with concurrency 1
Before
After (This PR)
Comparison with triton 3 stage
From #35

Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci