[NPU] Fix DeepEP LL dispatch BF16 flag and skip triton kernel on NPU for Qwen3.5 (#23815)
Co-authored-by: iridiumine <iridiumine@users.noreply.github.com>
This commit is contained in:
@@ -620,9 +620,13 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
|
|||||||
input_global_scale = self.quant_config.get("input_global_scale", None)
|
input_global_scale = self.quant_config.get("input_global_scale", None)
|
||||||
if input_global_scale is not None:
|
if input_global_scale is not None:
|
||||||
use_nvfp4 = True
|
use_nvfp4 = True
|
||||||
elif not get_moe_runner_backend().is_flashinfer_cutedsl():
|
elif not get_moe_runner_backend().is_flashinfer_cutedsl() and (
|
||||||
|
not _is_npu or not envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
|
||||||
|
):
|
||||||
# flashinfer_cutedsl expects BF16 dispatch when NVFP4 dispatch is
|
# flashinfer_cutedsl expects BF16 dispatch when NVFP4 dispatch is
|
||||||
# off; its kernel quantizes to NVFP4 internally.
|
# off; its kernel quantizes to NVFP4 internally.
|
||||||
|
# SGLANG_DEEPEP_BF16_DISPATCH forces BF16 dispatch for NPU
|
||||||
|
# where INT8 input + BF16 weight GMM is not supported.
|
||||||
use_fp8 = True
|
use_fp8 = True
|
||||||
|
|
||||||
# round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP
|
# round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP
|
||||||
|
|||||||
@@ -464,7 +464,11 @@ class Qwen3_5GatedDeltaNet(nn.Module):
|
|||||||
hidden_states
|
hidden_states
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.num_v_heads // self.num_k_heads in [1, 2, 4] and not _is_cpu:
|
if (
|
||||||
|
self.num_v_heads // self.num_k_heads in [1, 2, 4]
|
||||||
|
and not _is_cpu
|
||||||
|
and not _is_npu
|
||||||
|
):
|
||||||
mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous(
|
mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous(
|
||||||
projected_states_qkvz,
|
projected_states_qkvz,
|
||||||
projected_states_ba,
|
projected_states_ba,
|
||||||
@@ -488,6 +492,8 @@ class Qwen3_5GatedDeltaNet(nn.Module):
|
|||||||
query, key, value, z, b, a = self.fix_query_key_value_ordering(
|
query, key, value, z, b, a = self.fix_query_key_value_ordering(
|
||||||
projected_states_qkvz, projected_states_ba
|
projected_states_qkvz, projected_states_ba
|
||||||
)
|
)
|
||||||
|
b = b.contiguous()
|
||||||
|
a = a.contiguous()
|
||||||
|
|
||||||
query, key, value = map(
|
query, key, value = map(
|
||||||
lambda x: x.reshape(x.shape[0], -1), (query, key, value)
|
lambda x: x.reshape(x.shape[0], -1), (query, key, value)
|
||||||
|
|||||||
Reference in New Issue
Block a user