Conversation
blueswhen
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for SM100 GPU quantization (specifically using UE8M0 scales) across various Triton kernels, DeepGEMM implementations, and weight loading processes. It also optimizes memory usage by performing in-place weight transformations for Mega MoE instead of caching copies. The review feedback identifies a critical bug in transform_mega_moe_weights_in_place where w2.weight is not copied back, which would cause correctness issues or crashes. Additionally, a refactoring opportunity was identified to eliminate duplicated allocation logic for x_s in per_token_group_quant_fp8.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| w13.weight.copy_(transformed_l1[0]) | ||
| w13.weight_scale.copy_(transformed_l1[1]) | ||
| w2.weight_scale.copy_(transformed_l2[1]) |
There was a problem hiding this comment.
In transform_mega_moe_weights_in_place, the transformed weight for w2 (transformed_l2[0]) is not copied back to w2.weight. Only w2.weight_scale is copied. This will leave w2.weight in its original layout, which is incompatible with the Mega MoE kernel and will lead to correctness issues or runtime crashes.
| w13.weight.copy_(transformed_l1[0]) | |
| w13.weight_scale.copy_(transformed_l1[1]) | |
| w2.weight_scale.copy_(transformed_l2[1]) | |
| w13.weight.copy_(transformed_l1[0]) | |
| w13.weight_scale.copy_(transformed_l1[1]) | |
| w2.weight.copy_(transformed_l2[0]) | |
| w2.weight_scale.copy_(transformed_l2[1]) |
| if column_major_scales: | ||
| if scale_tma_aligned: | ||
| aligned_size = (x.shape[-2] + 3) // 4 * 4 | ||
| x_s = alloc_func( | ||
| x.shape[:-2] + (x.shape[-1] // group_size, aligned_size), | ||
| device=x.device, | ||
| dtype=torch.float32, | ||
| ).permute(-1, -2)[: x.shape[-2], :] | ||
| else: | ||
| x_s = alloc_func( | ||
| (x.shape[-1] // group_size,) + x.shape[:-1], | ||
| device=x.device, | ||
| dtype=torch.float32, | ||
| ).permute(-1, -2) | ||
| else: | ||
| x_s = alloc_func( | ||
| x.shape[:-1] + (x.shape[-1] // group_size,), | ||
| device=x.device, | ||
| dtype=torch.float32, | ||
| ) | ||
| lightllm_per_token_group_quant_fp8( | ||
| x, | ||
| group_size, | ||
| x_q, | ||
| x_s, | ||
| eps=eps, | ||
| dtype=dtype, | ||
| use_ue8m0_scales=use_ue8m0_scales, | ||
| ) |
There was a problem hiding this comment.
The allocation logic for x_s is completely duplicated between the if HAS_SGL_KERNEL and not use_ue8m0_scales: block (lines 146-166) and the else: block (lines 172-191). We can refactor this by moving the x_s allocation logic outside of the conditional block, which simplifies the code and improves maintainability.
Example refactoring:
if column_major_scales:
if scale_tma_aligned:
aligned_size = (x.shape[-2] + 3) // 4 * 4
x_s = alloc_func(
x.shape[:-2] + (x.shape[-1] // group_size, aligned_size),
device=x.device,
dtype=torch.float32,
).permute(-1, -2)[: x.shape[-2], :]
else:
x_s = alloc_func(
(x.shape[-1] // group_size,) + x.shape[:-1],
device=x.device,
dtype=torch.float32,
).permute(-1, -2)
else:
x_s = alloc_func(
x.shape[:-1] + (x.shape[-1] // group_size,),
device=x.device,
dtype=torch.float32,
)
if HAS_SGL_KERNEL and not use_ue8m0_scales:
finfo = torch.finfo(dtype)
fp8_max, fp8_min = finfo.max, finfo.min
sgl_ops.sgl_per_token_group_quant_fp8(x, x_q, x_s, group_size, 1e-10, fp8_min, fp8_max, False, enable_v2=True)
else:
lightllm_per_token_group_quant_fp8(
x,
group_size,
x_q,
x_s,
eps=eps,
dtype=dtype,
use_ue8m0_scales=use_ue8m0_scales,
)
1 participant
No description provided.