Skip to content

feat(rocm): integrate packed GR read for Qwen3.8 Flash Next TP2 and TP4 - #32

Draft
apinge wants to merge 1 commit into
qwen38-flash-next-amd-current-fix-installfrom
qwen38-flydsl-gr-read
Draft

apinge wants to merge 1 commit into
qwen38-flash-next-amd-current-fix-installfrom
qwen38-flydsl-gr-read

Conversation

@apinge

@apinge apinge commented Sep 18, 2026

Copy link
Copy Markdown
Owner

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

  1. in packed_weights apply preshuffling to both down gemm and up gemm weights via aiter.ops.shuffle.shuffle_weight For 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

 python -m sglang.test.run_eval \
  --eval-name gsm8k --num-examples 1319 --max-tokens 16384 --port 9080
[aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
[aiter] import [module_aiter_core] under /opt/aiter/aiter/jit/module_aiter_core.so
Downloading from https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl to /tmp/test.jsonl
/tmp/test.jsonl: 732kB [00:00, 3.87MB/s]                                                                                                                                                                            
ChatCompletionSampler initialized with self.system_message=None self.temperature=0.0 self.max_tokens=16384 self.reasoning_effort=None self.extra_body=None self.stop=None self.record_meta_info=True
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [21:23<00:00,  1.02it/s]
Total latency: 1283.525 s
Score: 0.979
Output throughput: 490.490 token/s
[METRIC] gsm8k_score=0.9794520547945206 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8", "eval": "gsm8k"}
[METRIC] gsm8k_latency=1283.5254440808203 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8", "eval": "gsm8k"}

Speed Tests and Profiling

profiled with concurrency 1

Before

image

After (This PR)

image

Comparison with triton 3 stage

From #35
image

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

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.
@apinge apinge changed the title feat(rocm): integrate packed GR read for Qwen TP2 and TP4 feat(rocm): integrate packed GR read for Qwen3.8 Flash Next TP2 and TP4 Sep 18, 2026
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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's remove markdown after PR ready. SGLANG_GR_READ_FLYDSL=1 should be mentioned in PR for enable this opt.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants