[Qwen4Exp] Reinterpret FP8 E4M3 PLE as int8 for SM86 pinned-host lookup (PR #58856)

The pinned-host PLE lookup Triton kernel only gathers and stores bytes,
but passing an fp8e4nv tensor to it fails on architectures without
native FP8 pointer support (e.g. SM86) with an error saying type
fp8e4nv is not supported (supported fp8 dtypes are fp8e4b15, fp8e5).

Since FP8 E4M3 and int8 both occupy one byte per element, reinterpret
the PLE weight view and the output tensor as int8 when the weight dtype
is float8_e4m3fn; the non-FP8 path is unchanged.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
2026-09-27 09:54:05 +07:00
co-authored by Hinata
parent 758cc4f9e1
commit 835c44af15
@@ -465,10 +465,17 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
flat_ids = input_ids.reshape(-1).long() flat_ids = input_ids.reshape(-1).long()
if flat_ids.numel(): if flat_ids.numel():
if self.weight.dtype == torch.float8_e4m3fn:
uva_weight_view = self._uva_weight.view(torch.int8)
output_view = output.view(torch.int8)
else:
uva_weight_view = self._uva_weight
output_view = output
_lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)]( _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
self._uva_weight, uva_weight_view,
flat_ids, flat_ids,
output, output_view,
self.embedding_dim, self.embedding_dim,
self.shard_indices.org_vocab_start_index, self.shard_indices.org_vocab_start_index,
self.shard_indices.org_vocab_end_index, self.shard_indices.org_vocab_end_index,