[Bugfix][Engram] Fix intermittent Triton 3.8 crash in the lookup kernel (#59639)
Signed-off-by: Djordje Ramic <djoramic@amd.com> Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
djramic committed
Oct 1, 2026 at 21:38 UTC
840c7c2f4bc72c7c55edb2dea6d23f9a9c9734e0
1 file changed
+11
-2
vllm/models/deepseek_v41/common/engram.py
+11
-2
@@ -605,6 +605,15 @@ def _engram_head_shard_weight_loader(
605
param.data.copy_(shard)
606
607
608
+# Branching on SORTED at runtime intermittently crashes Triton 3.8's
609
+# RemoveLayoutConversions.
610
+# TODO: Remove the 3.8 check once Triton ships
611
+# https://github.com/triton-lang/triton/pull/10706.
612
+_SORTED_IS_CONSTEXPR = tl.constexpr(
613
+ triton.__version__.startswith("3.8") and current_platform.is_rocm()
614
+)
615
+
616
+
617
@triton.jit(
618
do_not_specialize=[
619
"vocab_start",
@@ -613,7 +622,7 @@ def _engram_head_shard_weight_loader(
622
"ids_stride_t",
623
"ids_stride_h",
624
"GRID",
616
- "SORTED",
625
+ *(() if _SORTED_IS_CONSTEXPR else ("SORTED",)),
626
]
627
)
628
def _engram_lookup_kernel(
@@ -634,7 +643,7 @@ def _engram_lookup_kernel(
643
QUANT_BLOCK: tl.constexpr,
644
BLOCK_R: tl.constexpr,
645
GRID,
637
- SORTED,
646
+ SORTED: tl.constexpr if _SORTED_IS_CONSTEXPR else None, # type: ignore[valid-type]
647
):
648
"""Gather fp8 rows, apply their ue8m0 block scales, write bf16.
649