flashlib's package metadata scopes this project to NVIDIA ("Triton + CuteDSL kernels for NVIDIA GPUs", requires nvidia-cutlass-dsl), so this may simply be out of scope — filing for visibility since the kernel itself (kernels/slot_cache/triton/lru_ensure.py) is pure Triton with no NVIDIA-only intrinsics I could find, and it does execute (compiles, launches, doesn't crash) on AMD via Triton's ROCm backend.
Context: flashlib.kernels.slot_cache.lru_ensure is vendored by FlashML-org/FreeToken as the GPU-side LRU admission for its MoE expert offload cache (freetoken/moe/offload_kernels.py::ensure_experts). Running FreeToken on a discrete gfx1011 card (AMD Radeon Pro V520 / BC-160, Navi 12, ROCm via AMD TheRock nightly, torch==2.12.0+rocm10.2.0a20260921) serving a real GGUF MoE checkpoint (Qwen3.6-35B-A3B), decode-time output is corrupted (NaN or exact-zero logits depending on the surrounding config) while prefill is consistently clean.
What's been ruled out on the FreeToken side (full detail in FlashML-org/FreeToken#260, follow-up comment):
- Not the model's GDN/Gated-DeltaNet recurrence (swapped for a pure-PyTorch reference, same corruption).
- Not the native HIP GGUF/MoE compute kernels (forcing CPU-side expert compute still corrupts, differently).
- Not a kernel-launch-ordering race (
HIP_LAUNCH_BLOCKING=1 AMD_SERIALIZE_KERNEL=3, i.e. full serialization, does not fix it).
- Narrowed to: router and shared-expert outputs are clean at the exact forward call where the offload-cached routed-expert output (fed by
ensure_experts/lru_ensure) is NaN. Happens only on single-token decode, at a different transformer layer each run (consistent with depending on which specific expert ids collide for cache slots that step, not a fixed line).
- Checked the one hazard your own comment in
_lru_ensure_kernel flags explicitly (the CTA-scope fence guarding the post-_phase1 LRU-usage reload against a stale pre-bump view, tl.debug_barrier() at source line 184): pulled the compiled AMDGPU ISA for this kernel from the local Triton cache (gfx1011) and confirmed it lowers to a full s_waitcnt lgkmcnt(0) + s_waitcnt_vscnt null, 0x0 + s_barrier + buffer_gl0_inv sequence — correctly formed, not missing. So that specific documented hazard isn't it, at least not at that call site.
Not done / where I stopped: haven't audited the rest of _phase1/_install/the packed-key eviction math for other AMD-specific hazards (e.g. tl.argmin/atomic ordering, or the _insert streaming strategy's tiling, though decode's small K here should route to _seq not _insert), and haven't dumped the actual evict_slots/src_indices values at the moment of corruption (only NaN/absmax counters so far) — that would need either much more targeted instrumentation or interactive rocgdb. Filing this as reproducible evidence rather than a fix, since I don't know this codebase and this may well be a "not a supported target" non-issue for you — closing without action is completely reasonable if so, just wanted the repro on record in case it's useful (e.g. if the reload-fence class of hazard has a sibling elsewhere in the file that isn't as obviously commented).
Happy to share the raw .amdgcn/logs if useful.
flashlib's package metadata scopes this project to NVIDIA ("Triton + CuteDSL kernels for NVIDIA GPUs", requiresnvidia-cutlass-dsl), so this may simply be out of scope — filing for visibility since the kernel itself (kernels/slot_cache/triton/lru_ensure.py) is pure Triton with no NVIDIA-only intrinsics I could find, and it does execute (compiles, launches, doesn't crash) on AMD via Triton's ROCm backend.Context:
flashlib.kernels.slot_cache.lru_ensureis vendored byFlashML-org/FreeTokenas the GPU-side LRU admission for its MoE expert offload cache (freetoken/moe/offload_kernels.py::ensure_experts). Running FreeToken on a discrete gfx1011 card (AMD Radeon Pro V520 / BC-160, Navi 12, ROCm via AMD TheRock nightly,torch==2.12.0+rocm10.2.0a20260921) serving a real GGUF MoE checkpoint (Qwen3.6-35B-A3B), decode-time output is corrupted (NaN or exact-zero logits depending on the surrounding config) while prefill is consistently clean.What's been ruled out on the FreeToken side (full detail in FlashML-org/FreeToken#260, follow-up comment):
HIP_LAUNCH_BLOCKING=1 AMD_SERIALIZE_KERNEL=3, i.e. full serialization, does not fix it).ensure_experts/lru_ensure) is NaN. Happens only on single-token decode, at a different transformer layer each run (consistent with depending on which specific expert ids collide for cache slots that step, not a fixed line)._lru_ensure_kernelflags explicitly (the CTA-scope fence guarding the post-_phase1LRU-usage reload against a stale pre-bump view,tl.debug_barrier()at source line 184): pulled the compiled AMDGPU ISA for this kernel from the local Triton cache (gfx1011) and confirmed it lowers to a fulls_waitcnt lgkmcnt(0)+s_waitcnt_vscnt null, 0x0+s_barrier+buffer_gl0_invsequence — correctly formed, not missing. So that specific documented hazard isn't it, at least not at that call site.Not done / where I stopped: haven't audited the rest of
_phase1/_install/the packed-key eviction math for other AMD-specific hazards (e.g.tl.argmin/atomic ordering, or the_insertstreaming strategy's tiling, though decode's smallKhere should route to_seqnot_insert), and haven't dumped the actualevict_slots/src_indicesvalues at the moment of corruption (only NaN/absmax counters so far) — that would need either much more targeted instrumentation or interactiverocgdb. Filing this as reproducible evidence rather than a fix, since I don't know this codebase and this may well be a "not a supported target" non-issue for you — closing without action is completely reasonable if so, just wanted the repro on record in case it's useful (e.g. if the reload-fence class of hazard has a sibling elsewhere in the file that isn't as obviously commented).Happy to share the raw
.amdgcn/logs if useful.