[NVIDIA] Fix wrong symmetric sizes for fp4 cases (#12640)

This commit is contained in:
Kaixi Hou
2025-11-04 14:19:37 -08:00
committed by GitHub
parent 1cfbbc42d8
commit 34f7564df0

View File

@@ -1602,8 +1602,13 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
x_sf = nvfp4_block_scale_interleave(x_sf)
with use_symmetric_memory(get_tp_group()) as sm:
# The x might be packed in the case of fp4. So, use the output dim of the
# weight of the second GEMM.
symm_output = torch.empty(
x.shape[0], x.shape[1], dtype=output_dtype, device=x.device
x.shape[0],
layer.w2_weight.shape[1],
dtype=output_dtype,
device=x.device,
)
sm.tag(symm_output)