Skip to content

fix(casting): guard ANE topk and fix multi-output container downcasting - #178

Open
rohith500 wants to merge 1 commit into
apple:mainfrom
rohith500:fix/casting-topk-ane-guard
Open

rohith500 wants to merge 1 commit into
apple:mainfrom
rohith500:fix/casting-topk-ane-guard

Conversation

@rohith500

@rohith500 rohith500 commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #138

Two fixes for _FP16Casting:

  1. Guards torch.topk on ANE: Running aten.topk in FP16 on Apple Neural Engine packs two 16-bit indices into each 32-bit output word (Apple Radar FB24904516 / FP16 model silently returns wrong outputs with default specialization on iPhone 17 Pro; CPU-only, GPU-preferred and FP32 are correct coreai-torch#115), corrupting downstream indexing in detection models like RF-DETR and YOLO. Added guard_ane_topk: bool = True (with opt-out) to keep topk in FP32 with boundary casts.

  2. Downcasts multi-output containers at operator.getitem: Multi-output operations returning tuples (topk, split, chunk) that ran in FP32 skipped downcast insertion because container tuples lack .dtype. Extracted FP32 tensors leaked into downstream layers, forcing them into FP32 and triggering runtime weight upcasts. We insert downcasts on extracted float elements in operator.getitem (preserving FP32 if the activation overflowed during calibration or was explicitly ignored).

Tested with pytest tests/casting/ -v (142 passed) and make check (all 25 checks passed).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Multi-output container FP32 leak in _FP16Casting and ANE FP16 topk index packing erratum

2 participants