[NVIDIA] Fix wrong symmetric sizes for fp4 cases (#12640)
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user