fixes

Signed-off-by: Bill Nell <bnell@redhat.com>

Author: bnellnm

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -150,7 +150,9 @@ def maybe_make_prepare_finalize(
  hidden_dim_scale=hidden_dim_scale,
  in_dtype=in_dtype,
  out_dtype=in_dtype,
- scale_dtype=torch.float32,
+ scale_dtype=torch.float32
+ if quant_config.quant_dtype is not None
+ else None,
  max_private_tokens=None, # For tuning
  )
 

vllm/model_executor/layers/fused_moe/pplx_garden_prepare_finalize.py

@@ -25,7 +25,7 @@ def pplx_garden_hidden_dim_scale(
  quant_dtype: torch.dtype | str | None,
  per_act_token_quant: bool,
  block_shape: list[int] | None,
-) -> int:
+) -> int | None:
  # For blocked per token: set to
  # ceil_div(hidden_dim, block_size) * sizeof(float32)
  # For per-token: set to 4 * sizeof(float32) (x4 for alignment)
@@ -37,16 +37,16 @@ def pplx_garden_hidden_dim_scale(
  if per_act_token_quant:
  # per-token (M x 1)
  assert block_shape is None
- hidden_dim_scale = 1
+ hidden_dim_scale = 16
  elif block_shape is not None:
  # per-group (M x K_tiles)
  block_size = block_shape[1]
  hidden_dim_scale = cdiv(hidden_dim, block_size)
  else:
  # per-tensor (1 x 1)
- hidden_dim_scale = 1
+ hidden_dim_scale = 16
  else:
- hidden_dim_scale = 0
+ hidden_dim_scale = None # 1?
 
  return hidden_dim_scale
 
@@ -190,7 +190,7 @@ def prepare_async(
  expert_x_scale_shape = (
  self.num_local_experts,
  expert_x.size(1),
- round_up(final_dim, 4), # round up for alignment
+ round_up(final_dim, 16), # round up for alignment
  )
 
  expert_x_scale = torch.empty(
@@ -203,7 +203,11 @@ def prepare_async(
  # There's not much point setting this unless it is != indices.size(0)
  bound_m: torch.Tensor | None = None
 
- logger.debug("PPLX_GARDEN dispatch send %s", expert_x.shape)
+ logger.debug(
+ "PPLX_GARDEN dispatch send %s, %s",
+ expert_x.shape,
+ expert_x_scale.shape if expert_x_scale is not None else None,
+ )
 
  self.a2a.dispatch(
  out_expert_num_tokens=expert_num_tokens,
@@ -269,7 +273,8 @@ def _receiver(
  "PPLX_GARDEN receive X_SCALE %s",
  expert_x_scale.shape if expert_x_scale is not None else None,
  )
- logger.debug("PPLX_GARDEN receive META %s", expert_tokens_meta)
+ logger.debug("PPLX_GARDEN receive num_tokens %s", expert_num_tokens.shape)
+ # logger.debug("PPLX_GARDEN receive META %s", expert_tokens_meta)
 
  return expert_x, expert_x_scale, expert_tokens_meta, None, None
 
@@ -332,11 +337,13 @@ def finalize_async(
 
  logger.debug("PPLX_GARDEN combine send")
 
+ hidden_dim = output.size(1)
+
  self.a2a.combine(
  out_tokens=output,
  indices=topk_ids_u32,
  weights=topk_weights,
- expert_y=fused_expert_output,
+ expert_y=fused_expert_output.view(-1, hidden_dim),
  bound_m=bound_m,
  do_send=True,
  do_recv=False,
@@ -349,7 +356,7 @@ def finalize_async(
  out_tokens=output,
  indices=topk_ids_u32,
  weights=topk_weights,
- expert_y=fused_expert_output,
+ expert_y=fused_expert_output.view(-1, hidden_dim),
  bound_m=bound_m,
  do_send=False,
  do_recv=True,

vllm/model_executor/layers/quantization/fp8.py

@@ -1040,6 +1040,9 @@ def select_gemm_impl(
  BatchedTritonOrDeepGemmExperts,
  TritonOrDeepGemmExperts,
  )
+ # from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
+ # NaiveBatchedExperts,
+ # )
 
  assert not self.use_marlin and not self.rocm_aiter_moe_enabled, (
  "Marlin and ROCm AITER are not supported with all2all yet."
@@ -1061,6 +1064,11 @@ def select_gemm_impl(
  self.weight_block_size,
  False,
  )
+ # return NaiveBatchedExperts(
+ # max_num_tokens=max_num_tokens_per_rank,
+ # num_dispatchers=prepare_finalize.num_dispatchers(),
+ # quant_config=self.moe_quant_config,
+ # )
  return BatchedTritonOrDeepGemmExperts(
  max_num_tokens=max_num_tokens_per_rank,
  num_dispatchers=prepare_finalize.num_dispatchers(),