47 lines
1.9 KiB
Diff
47 lines
1.9 KiB
Diff
--- a/vllm/_custom_ops.py
|
|
+++ b/vllm/_custom_ops.py
|
|
@@ -1933,14 +1933,28 @@
|
|
if scale is None:
|
|
if use_per_token_if_dynamic:
|
|
scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32)
|
|
- torch.ops._C.dynamic_per_token_scaled_fp8_quant(
|
|
- output, input, scale, scale_ub
|
|
- )
|
|
+ # COREXX_FP8_QUANT_PATCH: native fallback (vLLM _C op missing on CoreX)
|
|
+ m = input.shape[0]
|
|
+ s = input.abs().amax(dim=1, keepdim=True).float().clamp_min(1e-12) / 448.0
|
|
+ if scale_ub is not None:
|
|
+ ub = torch.as_tensor(scale_ub, device=s.device, dtype=s.dtype)
|
|
+ s = torch.minimum(s, ub)
|
|
+ scale[:m].copy_(s)
|
|
+ if shape[0] > m:
|
|
+ scale[m:].fill_(1.0)
|
|
+ output[m:].zero_()
|
|
+ output[:m].copy_((input.float() / s).to(out_dtype))
|
|
else:
|
|
scale = torch.empty(1, device=input.device, dtype=torch.float32)
|
|
- torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
|
|
+ # COREXX_FP8_QUANT_PATCH: native fallback
|
|
+ amax = input.abs().amax().float().clamp_min(1e-12)
|
|
+ scale.fill_(amax / 448.0)
|
|
+ output.copy_((input.float() / scale).to(out_dtype))
|
|
else:
|
|
- torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape)
|
|
+ # COREXX_FP8_QUANT_PATCH: native fallback
|
|
+ if group_shape is not None and tuple(group_shape) not in ((-1, -1), (-1, 1), (1, -1)):
|
|
+ raise NotImplementedError(f"native fp8 static quant: unsupported group_shape {group_shape}")
|
|
+ output.copy_((input.float() / scale).to(out_dtype))
|
|
|
|
return output, scale
|
|
|
|
@@ -3824,6 +3838,9 @@
|
|
mul_routed_weight,
|
|
top_k,
|
|
config['BLOCK_SIZE_M'],
|
|
+ # COREXX_FP8_MOE_PATCH: pass quant scales to ixformer kernel
|
|
+ w_scale=B_scale,
|
|
+ a_scale=A_scale,
|
|
bias=bias
|
|
)
|
|
|