Skip to content

[CPU]convert topk_weights to fp32 for INT8 and FP8 paths (for llama4) and fix LmHead weight pack #7818

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhyncs merged 4 commits into from
Jul 9, 2025
Merged

[CPU]convert topk_weights to fp32 for INT8 and FP8 paths (for llama4) and fix LmHead weight pack #7818

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6d09e3e
[CPU] convert topk_weights to fp32 for INT8 and FP8 paths
chunyuan-w Jul 7, 2025
0655df1
Merge branch 'main' into chunyuan/fix_int8_fp8_llama4
zhyncs Jul 7, 2025
4a9087f
convert topk_weights to FP32 in kernel
chunyuan-w Jul 8, 2025
d92042a
For LmHead weight pack, check the weight dtype explicitly
chunyuan-w Jul 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,7 +319,7 @@ def forward_cpu(
layer.w2_weight,
topk_weights.to(
torch.float
), # TODO: the topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bfloat16 while the kernel requires it to be float32
), # TODO: [Note] topk_weights should be float32 for now. The topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bfloat16 while the kernel requires it to be float32
topk_ids,
False, # inplace # See [Note] inplace should be False in fused_experts.
False, # use_int8_w8a8
Expand Down
4 changes: 3 additions & 1 deletion python/sglang/srt/layers/quantization/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1002,7 +1002,9 @@ def apply(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_weights.to(
torch.float
), # See [Note] topk_weights should be float32 for now.
topk_ids,
False, # inplace See [Note] inplace should be False in fused_experts.
False, # use_int8_w8a8
Expand Down
4 changes: 3 additions & 1 deletion python/sglang/srt/layers/quantization/w8a8_int8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,9 @@ def apply(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_weights.to(
torch.float
), # See [Note] topk_weights should be float32 for now.
topk_ids,
False, # inplace See [Note] inplace should be False in fused_experts.
True, # use_int8_w8a8
Expand Down