Skip to content

[Refactor] Unify indirect-access offset paths, add example tests #212

Description

@yuhaohaoyu

Scope/Background

PR #188 (and its predecessors #147, #183, #193) delivered a stable blocked-indirect fast path for emulating kernels aligning with MoE, paged-attention, and multi-index gather/scatter patterns. While the offsets= refactor unified read/write steps, two independent offset-computation paths remain: the general (item-wise) path (_resolve_idx_reads_build_indirect_coords_flat_memory_offsets) and the fast path (_prepare_dep_var_sub_space_runtime_read_and_expand_sub_space_gen_offsets_vso_space_via_broadcast), where both produce the same artifact — a flat int64 offset array — via completely separate code.

Problem to Address (according to Lars' final comments fro PR 188)

  • **M1 (dual paths):**The dual-path existence hurts maintainability. The shared-view (shared subscript array) fix in _build_indirect_coords (PR Blocked-indirection load store, fast path #188) is a concrete example for reference when addressing the issue.
  • Test gap: No existing MLIR example exercises the fast path end-to-end. All indirect-access fixtures (including paged-attention/paged-tensor) fail the has_direct_expr if-gate, making the fast path unreachable from example tests. The fix is needed for the execution of notebooks, etc.

Scope/Proposal

  • Refactor both pipelines (element-wise and blocked-index) to share a common offset-computation interface, where the fast path is a vectorized specialization of the general path rather than a parallel implementation. Lars's recommendations file outlines the structural approach.
  • Add 2–3 MLIR integration fixtures (MoE-style, paged-attn-style) that satisfy all fast-path gates, confirming the blocked-indirect path is exercised by pytest without synthetic Python-only test scaffolding.

Acceptance Criteria

  • Single offset-computation interface (internal) for the element-wise and blocked-index emulation paths for indirect-accesses.
  • Example tests that emulate MoE and PagedAttn kernels with the blocked-index emulation path lit.

Out of Scope

Activity

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

Metadata

Metadata

Assignees

Labels

bugSomething isn't workingdocumentationImprovements or additions to documentationgood first issueGood for newcomers

Projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions