Blocked-indirection load store, fast path

[yuhaohaoyu]

Purpose

Thorough rework of the blocked indirection fast emulation proposal.

Work title: Blocked-indirect fast path: read K indices, broadcast to N offsets

• Factor indirection dims from static dims — read K values, meshgrid to N offsets in one vectorized pass. 100–190× emulation speedup for MoE/paged-attn patterns. fixes [Bug] block-gather fast path (#147) lost zero-padding for reads past an allocation — partial tile loads now raise IndexError #182 zero-padding bug.

Related issues, PR

1. Round 1: issue [Feature] Fast Indirection Emulation targeting Blocked Indirection #143, PR Add block-gather fast path for indirect memory access #147 : merged after extensive review and latter reverted for clean refactoring planning reached with team agreement.
2. Round 2: issue [refactor] improved integration of block-gather fast path #171, PR Issue 171, part 1, follow up of pr-147 on block-indirection emulation #179 : quick refactoring before the reversion of PR 147.
3. Round 3: issue [Bug] block-gather fast path (#147) lost zero-padding for reads past an allocation — partial tile loads now raise IndexError #182 zero-padding bug.

Scope of this PR

1. Complete addressing the issue 171 and 182.
2. Merge of the PR should lead to PR-179 closing without merging: Treat it as through re-lifting, with loyal referencing to review comments in PR-147 (code+), PR-179 (code-).

[yuhaohaoyu]

(Manually written, based on work-session notes).

Note for reviewing the feature: Blocked-Indirection

The key algorithmic idea: for loads like A[B[e], m, n], instead of computing addresses for all N points, we read K index values (the indirection dimension) and broadcast them across the static dimensions using NumPy meshgrid. This specialization treatment, together with numpy gather/scatter ops, yielding 100–190× emulation speedup. This note documents the design choices, performance justification for keeping gather/scatter.

Introduction

Indirect-access in array-centric language resembles the pointer-based accesses in C. Given its complexity, pattern based specialization are beneficial for speed of emulator's runtime.

Broadly, data mapping between spaces can be considered as reducing/collecting/selecting (large to small spaces), permuting (similar sized spaces), scatter/replicating/broadcasting (small to large spaces). Indirect-access patterns from LLM are mostly fall in large to small spaces or between similar sized spaces. Utilizing the gather vectorized ops from Numpy for realistic sizes of data is natural to boost emulation speed.

Given that LLM kernels are native to GPU or matrix-engines that are fluent with dense metrics, arrays and vectors (contrasting to linked-list, graph-traversal and other irregular access patterns), there are not many indirect-carrying kernels appearing in popular LLMs and many of them carry static inner dimensions. Such pattern (the indirection is not in all the subscript expressions) of a KTIR load statement need not to compute indices/addresses of the source space for all the elements under read. Instead, indices of elements (in the source space) in the same static dimension can be produced through Nympy's broadcast op (similar to everyday replication, or splat in SIMDization).

For ease of reviewing, many algorithmic, logic considerations are kept in docstrings and comments. Following sections are materials and points useful for the decision making on choices of coding, design, reviewing, and maintenance.

Key notes for reviewing:

1. PR Issue 171, part 1, follow up of pr-147 on block-indirection emulation #179 was the shallow attempt to refactoring the blocked-indirect code path to be aligned with the general, element-wise, indirection path. Massive great review points were posted there. This deep refactoration complies 8 out 9 of the revision suggestions, only kept the 'gather/scatter' treatment in data loading.
2. Section 2 describes the code change sets, organized by exsecution stage.
3. Gather/Scatter are Numpy ops show 3x load-time reduction when compared with span-copy used in the general code path. Particularly when the data moving is in a gathering/collecting pattern and with static inner-most dimensions (MoE-FFN, Paged-attn)
4. Seciton 4 explains the come-and-go aliasing issue: introduced by _place_in_lx, cleared by revert back to use _write_in_lx.
5. The data flow of what the steps of the emulation of a tensor load with indirections.
6. What's the common patterns of indirect accesses in LLM that'd motivated the investigation and led to 150x emulation efficiency.
7. Section 7 lists a few complex patterns that the PR is partially or not covering.
8. Zero-padding error in indireciton load reported in issue [Bug] block-gather fast path (#147) lost zero-padding for reads past an allocation — partial tile loads now raise IndexError #182 is addressed in the PR.

1. Compliance with PR #179 Review Comments

#	Reviewer ask	Status	What we did
1	Delete _is_block_gather	Done	Removed; callers use _analyze_blocked_indirect + callers check the returned tuple directly
2	Delete _block_gather_read_idx, use shared idx-read helper	Done	_runtime_read_and_expand_sub_space shared by fast + general path
3	Delete HBMSimulator.gather/scatter, revert to _read_flat/_write_flat	Kept — justified	O(selected) vs O(span); sparse patterns lose 2-3× speedup without it. .
4	Delete _place_in_lx, revert to _write_to_lx	Done	All 4 call sites reverted; method deleted
5	Delete _block_gather_offsets_fallback	Done	dropped the fallback code path for irrealistic corner case
6	No PR numbers in function names	Done	Benchmark code trimmed to covering 2 paths, minimal in splashing
7	Revert uv.lock	Done
8	Benchmark: parameterized, not hardcoded recipes	Done	Removed cmd_4way, cmd_gather, subcommand dispatch. Single-purpose TOML-driven script.
9	indirect_store next to indirect_load	Done	Moved; now adjacent (L912, L962)

Summary: 8/9 items fully addressed, 1 justified divergence (gather/scatter kept for perf).

---

2. Change Set: Blocked-Indirect Path (Current State)

Entries and call chains

ktdp_ops.py:194  →  MemoryOps.indirect_load   ─┬─ if _analyze_blocked_indirect → fast path
                                                └─ else → general path

ktdp_ops.py:232  →  MemoryOps.indirect_store  ─┬─ if _analyze_blocked_indirect → fast path
                                                └─ else → general path

---

stages and func codes — memory_ops.py (ordered by execution)

Stage	Function	Size	Role
0. Helpers	_expr_dependent_vars	23	AST walk: extract iteration-dim dependencies
1. Classify	_analyze_blocked_indirect	78	Gates 1-3, returns analysis tuple or None
2. Orchestrate/Umbrella	_compute_blocked_indirect_offsets	25	Chains stages 2→3→4
2.1. Prepare subspace	_prepare_dep_var_sub_space	23	K sampling coordinates for dep-var subspace
2.2. Read & expand	_runtime_read_and_expand_sub_space	59	K index values per sub from HBM (scattered DMA)
2.3. Broadcast	_gen_offsets_vso_space_via_broadcast	72	K→N: meshgrid over dep+direct axes → N flat offsets
3. Entry (load)	indirect_load	48	Fast path via stage 5, or general path fallback
0. Entry (store)	indirect_store	49	Mirror of indirect_load for scatter
		377	Fast path total

What is a bit new (not in PR 147)

Addition	Purpose
Per-sub re-keying in _runtime_read_and_expand_sub_space	Accumulates per subscription expression, not per view — enables shared-view patterns - Section 7 on complex patterns

General path (shared infrastructure, also in memory_ops.py)

Function	Size	Role
_element_offsets	21	Reference impl: per-point Python loop
_enumerate_in_vso_order	17	VSO-permuted point enumeration
_resolve_idx_reads	17	General-path wrapper over stage 3
_build_indirect_coords	47	Per-point coord materialisation
	102	General path total

Tests

File	Lines	Tests	Category
test_blocked_indirect_fast_path.py	847	35	Fast path: gating, load, store, equivalence, shared-view, edge cases
test_indirect_access.py	889	21	Full stack: MLIR → IAT → indirect_load/store → HBM verification
	1736	56

Benchmark

File	Lines	Category
bench_indirect_emul_time.py	256	Single-purpose: element-wise vs block path, TOML-driven

Totals by category

Category	Lines
Fast-path implementation	377
General-path implementation	102
Tests	1736
Benchmark	256
Total	2471

Notes on: _expr_dependent_vars vs eval_subscript_expr

Function	Purpose	Returns
eval_subscript_expr(expr, pt)	Evaluate expr at one point	int — the computed value
_expr_dependent_vars(expr)	Static analysis: which vars does expr reference?	set — variable indices

_expr_dependent_vars is purely a block-path concept — it identifies which variables can be factored out (dep_vars → K reads) vs broadcast (direct dims → free amplification). The element-wise path doesn't need this analysis because it evaluates everything at every point anyway (brute-force N iterations).

They both walk the same AST structure, but for different purposes: one computes values, the other extracts metadata.

---

Rename map (PR 147 → current)

PR 147 name	Current name	Rationale
_block_gather_analyze	_analyze_blocked_indirect	Describes what it does, not the marketing name
_dep_var_points	_prepare_dep_var_sub_space	Clarifies role: defines the dep-var subspace
_read_idx_at_points	_runtime_read_and_expand_sub_space	Clarifies: runtime DMA read that fills the subspace values
_block_gather_offsets	_gen_offsets_vso_space_via_broadcast	Describes the K→N expansion over full VSO space
_get_blocked_indirect_offsets	_compute_blocked_indirect_offsets	Active verb, consistent naming
_block_gather_load	(inlined into indirect_load)	No wrapper needed
_block_gather_store	(inlined into indirect_store)	No wrapper needed

---

3. Performance Impact with and without using op gather

Workload	Current block time	Extra span copy	Block time if reverted	Speedup impact
sparse-attn-128K	1.3 ms	+3.2 ms (32MB copy)	~4.5 ms	175× → ~50×
sparse-attn-32K	0.5 ms	+0.4 ms (4MB copy)	~0.9 ms	114× → ~63×
moe-1M	9.4 ms	+6.4 ms (64MB copy)	~16 ms	117× → ~69×
multi-head-4M	40 ms	+6.4 ms (64MB copy)	~47 ms	190× → ~164×

The sparse-attention workloads get destroyed. The span is 64-128× larger than the data actually needed. Copying the entire parent allocation (a maybe target to treat for future) just to fancy-index a tiny subset is the exact inefficiency that gather eliminates.
Extra span copy column: the cost of reading the entire contiguous allocation into a temp buffer before indexing into it.

• Dense workloads (multi-head, paged-attn with ratio 8×): moderate impact (~15% regression of load time)
• Sparse workloads (sparse-attn with ratio 64-128×): speedup drops by 2-3× (from 175× to ~50× of load time)

gather/scatter is not just "slightly better" — for sparse access patterns, it's the difference between O(selected) and O(parent_size) data movement. The implementation is tiny, ~one line (existing.ravel()[offsets]) and serves the general strided/coords path too (not block-gather-specific duplication).

---

Hypothetical: Refactoring Gain if gather/scatter Removed (Ignoring Perf)

If we ignore the performance regression, removing gather/scatter would eliminate:

Deleted	Lines	Location
_gather_from	~3	memory.py
_scatter_into	~5	memory.py
HBMSimulator.gather/scatter	~12	memory.py
LXScratchpad.gather/scatter	~12	memory.py
_MemAccessor.gather/scatter	~20	memory_ops.py
Total	~52 lines

gather/scatter removal: ~52 lines saved, no conceptual simplification, API vs inline code replication, bad for maintainability.

---

4. lasch's Review Item #1 — _place_in_lx Aliasing

From lasch's PR 147 review:

_place_in_lx aliasing contract is violated — All three call sites produce a C-contiguous alias: data, LX.memory[lx_ptr], and Tile.data all point to the same backing buffer. linalg.matmul and linalg.batch_matmul both mutate acc.data in-place with +=, writing through the alias.

His recommendation: "change _place_in_lx to store data.ravel().copy(). Track as a follow-up issue."

Our A3 resolves this exactly — reverted all 4 _place_in_lx call sites to _write_to_lx, which goes through _write_flat → data.flatten() (copy). The aliasing is eliminated. Cost: ~2ms on 40ms block path (190× → 183×).

5. What load Actually Does — Data Flow

HBM dict (np.ndarray)  ──fancy-index──►  new array  ──_write_to_lx──►  LX dict
                         (one copy)        ("wire")    (second copy)

• data.ravel()[offsets] is the copy — allocates output array, fills with selected elements in one pass. No intermediate temp buffer.
• The "new array" is a wire — carries the result from the gather unit to the LX write port. No semantic meaning, no lifetime. A Python artifact modeling what hardware does in a single DMA shot.
• On real hardware: one DMA transfer (HBM → LX SRAM with gather unit selecting elements). The simulator models semantics, not the exact transfer mechanism.

---

6. Motivational patterns: MoE and Paged Attn

The patterns in LLM kernels critical for IBM and community suggests the need for the specialization treatment: majority of the indices into the source of the data load can be calculated in bunks, no need to instantiate element by element. A block-factor, in the table, of 8192× means the fast path reads 8192× fewer index values than the element-wise path would.

#	Test	Indirection pattern, and block factor
1	test_moe_pattern	X[IDX[e], m, n], 8 from 128, block_factor=8192x
2	test_paged_attention_pattern	cache[BT[0,d0], d1, d2, d3], 1 page, block_factor=16384x

---

7. Complex Workload Patterns partially treated

The view-overlap gate rejects any IAT where multiple indirect subs reference the
same index_view_idx. This blocks real workload patterns:

Pattern	Representation	Why rejected
A[B[i][0], B[i][1], B[i][2], n]	2D index table, different columns	3 subs share view B
A[B[e], B[e+1], m, n]	1D index array, shifted window	2 subs share view B
A[B[e], B[h], m, n]	1D index, different dep_vars	2 subs share view B
A[B[e], B[e], m, n]	Same expr (degenerate diagonal)	2 subs share view B

These patterns appear in: sparse tensor coordinate tables, multi-field block
tables (paged attention), shifted-window attention, multi-dimensional gather.

---

8. Issue #182: Block-Gather Zero-Padding Bug

1) Error emerging workload/test

Triton-generated variable-length flash-attention prefill kernel:

• Descriptor: 100 rows loaded through a 128-row access tile (fp16, shape [4, 100, 64], per-head stride 100*64=6400)
• Final head tile iterates rows 0–127, but only 0–99 exist in the allocation (25,600 elements total)
• Offsets for rows 100–127 compute to indices 25,600+ — past the end of the flat allocation
• Result: IndexError: index 25600 is out of bounds for axis 0 with size 25600

2) What the spec says about zero-padding

The KTIR spec (RFC 0682) is silent on OOB read semantics. ktdp.load is defined simply as "reads data using an access tile's coordinates; produces a tensor" — no mention of what happens when coordinates fall outside the allocation.

However, the emulator (_read_flat in memory.py:224–255) has an explicit zero-padding contract:

"Elements beyond the end of the containing allocation are zero-padded."

3) Root cause

Path	OOB handling	Mechanism
_read_flat (strided)	Zero-pads	Allocates np.zeros(n), copies only available elements
_gather_from (gather)	Raises IndexError	Direct data.ravel()[offsets], raw numpy fancy-index — no bounds check, no padding.

Block-gather computes offsets via numpy broadcast (_gen_offsets_vso_space_via_broadcast). When a partial tile has tail elements mapping past the allocation end, they just point past the physical allocation boundary.

General path pseudo code

    # target space allocated and zero-out with the extent (max size)
    load_target_space = alloc_zero_out( 128 * 64 )
    load_data_to_non_zero_locs(...) --> load_target_space

Block path pseudo code

    size_space_indirect_subscripts = 100
    subs = load_indirect_subs( size_space_indirect_subscripts )
    fancy_index_offsets = numpy.broadcast( subs, 64 )  
    assert sizeof( fancy_index_offsets ) == 100 * 64
    fancy_index_load_data( fancy_index_offsets ) --> load_target_space
    # target space size determined by numpy broadcast and allocated to fit the non-loaded data.

Fix (pseudo code)

_gather_from — clamp OOB (out of bound) indices, zero-fill:

flat = data.ravel()
indices = elem_offset + offsets
oob = indices >= flat.size
indices = np.where(oob, 0, indices)
result = flat[indices]
result[oob] = 0 # OOB slots get zero, matching the general path - _read_flat
return result

_scatter_into — silently drop OOB writes:

indices = elem_offset + offsets
inbounds = indices < buf.ravel().size
buf.ravel()[indices[inbounds]] = data[inbounds]

[lasch]

I had the AI rummage through the code and the history of this overall effort since 147 with all the comments and experiments taken into account. Not going to blast this essay here (yet) because it doesn't really reflect the essence (imho).

To me it boils down to 'keep it simple to improve maintenance' vs. 'speed it up to not become the bottleneck' (see the push for a Rust clone). Therefore:

Question: do the benchmarks stand on their own feet? I.e. Could they be extracted into their own PR?
If yes, I'd advice to:

• create a PR with the benchmarks first
• run the benchmarks on the experimental codes that were suggested here
• use the collected evidence to validate the shipped offsets= design in comparison to coords=-reuse design

If no: could there be an adjusted (temporary) version of the benchmarks to get numbers from the experimental codes?

[yuhaohaoyu]

@lasch great suggestions. acting on it now.

[lasch]

PR #193 has merged. The next steps for this one are now:

• collect benchmark data from the experimental branch
• remove the benchmark codes (except for some intentional updates to include any new code)

[yuhaohaoyu]

[This should have some iterations, I will just post PDF files from some md files.]

@lasch comments done.

benchmark data for: contrasting table of this branch against fabian's private branch

• in the PDF file. response-188-post-193-merge.md.pdf

• perf diffs:

  1. Dense patterns (moe, multi-head): 1.0–3.1x faster end-to-end;
  2. Sparse patterns (sparse-attn, retro-attn): 3.8–194x faster — O(selected) vs O(span) data movement

code rebase + cleanup benchmark
rebase done, yet not started work on clean yet.

[lasch]

Confirming benchmark results reproduced with roughly the same results. 👍
Will continue review after cleanup is complete.

[yuhaohaoyu]

@lasch The final cleanup and soul-searching to avoid surprise like issue #182 are done.

Attached is a cut-down version of the report on adding corner case tests. I found it's mostly good and useful.

• pr-188-final-review-addressing-short.md

Commit messages:

Address PR-188 review: remove superseded benchmarks, add correctness gap tests

• Remove bench_indirect_emul_time.py and indirect_emul.toml (superseded by PR Add implementation-agnostic benchmark for indirect_load #193)
• Add tests for non-zero vss.lo in fast path dep-var iteration (no bug)
• Add tests for negative index guard asymmetry between upstream and primitives
• Add xfail test for shared-view general-path crash

On the xfail test targeting an issue in the general-path crash

The detail is in the md file in section tiled as 'Gap 3'. Should be a separate issue. Yet it's quite a corner case.

[lasch]

Add xfail test for shared-view general-path crash

The detail is in the md file in section tiled as 'Gap 3'. Should be a separate issue. Yet it's quite a corner case.

The chain of thought is:

• the spec allows shared-index-views, so it's good to get started supporting it
• but they're marked as 'known gap' because the general path crashes (as you documented)
• the new fast path now does 2 things:
  1. it introduces the possibility to use shared-index-view
  2. only the fast-path is functional

-> This is introducing inconsistent behavior depending on the input (the _analyze_blocked_indirect gate). It means that the gate is not just deciding general vs. fast, it also decides crash vs. functional.
I understand that the most common/expected workloads won't hit the general path. I'd still prefer consistent behavior.

One thing we missed:

One proposal from @fabianlim was (and I rephrase):
"... add an offsets= parameter directly to _read_flat/_write_flat themselves, keep gather/scatter only as thin forwarding wrappers (or delete them and call the parameterized primitive directly)"

Claude's assessment of this option expects: preserves the exact same hot call and therefore the exact same performance property, while collapsing the primitive-family count from three (_read_flat/_write_flat, gather/scatter, and whatever a future third mode would add) to one. It would also naturally share OOB/zero-pad handling with _read_flat's pre-existing contract instead of maintaining two separate implementations of "what happens at the boundary" (_read_flat's zero-pad vs. _gather_from/_scatter_into's OOB clamp, currently two co-existing fixes for the analogous problem in two shapes of access)."

I think, this option has not been tried out or measured/evaluated yet, is that correct @yuhaohaoyu ?
it would reduce the weight of the raised Maintenance concerns.

I'm pasting the detail about this verbatim without summarization:

Addendum: a middle ground for M1/M2, given the G3 performance evidence

M1/M2 keep being flagged as design costs, but G3's data shows the gather/scatter primitive family is buying something real on sparse patterns. Whether there's a way to keep that performance and shed the design cost splits into two different questions depending on which of fabianlim's two #147/#179-era proposals (§ "The fabianlim #147 'suggested changes'," above) is on the table — and the newly-collected G3 evidence only speaks to one of them.

M2 (the parallel gather/scatter primitive family) has a candidate low-cost fix: proposal (a). Proposal (a) — from fabianlim's first #147 comment (§3.1), never implemented as a branch — is to add an offsets= parameter directly to _read_flat/_write_flat themselves, keeping gather/scatter only as thin forwarding wrappers (or dropping them and calling the parameterized primitive directly). The entire sparse-pattern win, per the step-breakdown data (the P1–P4 comparison, and the T_indirect_load/T_load/T_offsets breakdown from the post-#193 PDF), is localized to the terminal data-movement step — reading exactly the selected elements (data.ravel()[offsets], O(selected)) instead of a full contiguous span-then-slice (O(parent_size)). Nothing about that win depends on gather/scatter being a separate, third primitive family alongside _read_flat/_write_flat; (a) would issue the identical ravel()[offsets] call, just relocated into a branch of an existing function rather than a separately-named one. If it works as described, it collapses the primitive-family count from three to one and unifies OOB/zero-pad handling (_read_flat's existing zero-pad contract vs. _gather_from/_scatter_into's separate OOB clamp — currently two implementations of the same boundary problem for two shapes of access).

Proposal (b) is a different, more aggressive move, and is what the newly-collected G3 data actually measures. Proposal (b) is fabianlim's block-gather-reuse-validation branch (@ 6bb4d9e) — the one this update's T_indirect_load/T_load/T_offsets reproduction ran against the shipped offsets= design. It deletes gather/scatter entirely and routes everything through the general coordinate-set path (MemoryOps.load(coords=...) → _flat_memory_offsets, which expects List[Tuple[int, ...]]). As the original review's testing-gaps analysis worked out, that path's _build_indirect_coords materializes one Python tuple per point — the 17.2%-of-general-path-time cost that fabianlim's own proposal never resolved definitively (by his own admission, "I didnt check the benchmarks... could have missed the whole point of the optimization"). The reproduced numbers (sparse patterns losing 2–3× when data movement is forced through this route) confirm that risk was real, not hypothetical — a solid, independently-verified result.

The two proposals must not be conflated, and only one of them has been measured. Proposal (a) was never built into a branch — it exists only as prose in fabianlim's first comment — so nothing in this PR's benchmark history, including the newly-collected data, has ever measured it. The claim that (a) would preserve the shipped design's performance is an architectural inference (same hot call, no per-point coordinate-tuple materialization, so no obvious new cost) rather than a benchmarked result. It's a reasonably strong inference precisely because (a) looks like a relocation of existing code rather than a different algorithm — unlike (b), which genuinely changes the data path — but it remains a proposal to prototype and benchmark, not a conclusion the evidence collected so far already establishes. What the data does establish is narrower and still useful: (b)'s specific mechanism (coordinate-tuple materialization) has a real sparse-pattern cost, which is one good reason to prefer (a) over (b) if either is pursued — it just isn't yet a demonstration that (a) is free.

M1 (the two independent offset-computation pipelines: general per-point enumeration vs. fast K-dedup-then-broadcast) does not have an equally clean middle ground. The K-vs-N dedup is where the dominant 74.5%-of-general-path-time win actually lives (per the G3 step-breakdown analysis), and that's a structurally different mechanism from anything either of fabianlim's proposals touches — both (a) and (b) still read the reduced K-sized point set; they only differ in what happens after. Unifying _resolve_idx_reads/_build_indirect_coords with _prepare_dep_var_sub_space/_runtime_read_and_expand_sub_space/_gen_offsets_vso_space_via_broadcast into one parameterized, vectorized pipeline is the "close-formed refactor" yuhaohaoyu's own #147-era comment already promised as later work — genuinely harder, unbenchmarked by anyone so far, and not resolved by either proposal.

(checked past comments on issues and PRs for discussion about proposal (a) but didn't find any substantial counter-arguments)

[yuhaohaoyu]

@lasch Addressing for both points are coming in. The benchmarking was great. Here to post the no-regression-on-perf messaging:

bench_indirect_load_path_compare — indirect-load-path-timing
=================================================================================================================================================
Workload         | Access expr               | Source shape         | Result shape       | Model         | T_indirect_load |   T_load | T_offsets
-----------------+---------------------------+----------------------+--------------------+---------------+-----------------+----------+----------
moe-262K         | X[IDX[e], m, n]           | [128, 256, 128]      | [8, 256, 128]      | deepseek-v2   |         2.40 ms |  1.88 ms |   0.52 ms
moe-1M           | X[IDX[e], m, n]           | [128, 1024, 128]     | [8, 1024, 128]     | deepseek-v2   |        10.18 ms |  7.81 ms |   2.38 ms
paged-attn-256K  | cache[BT[d0], d1, d2, d3] | [1024, 16, 128, 128] | [16, 16, 128, 128] | llama-8b      |        33.24 ms | 28.00 ms |   5.25 ms
sparse-attn-32K  | X[P[e], T[e], h]          | [1024, 512, 64]      | [4, 16, 64]        | bigbird-base  |         0.17 ms |  0.03 ms |   0.14 ms
sparse-attn-128K | X[P[e], T[e], h]          | [4096, 512, 256]     | [4, 16, 256]       | bigbird-large |         0.26 ms |  0.11 ms |   0.16 ms
retro-attn-16K   | mem[doc[d], chunk[c], h]  | [1024, 128, 128]     | [32, 4, 128]       | retro-7b      |         0.41 ms |  0.15 ms |   0.26 ms
multi-head-4M    | X[E[e], H[h], m, n]       | [128, 64, 256, 256]  | [8, 8, 256, 256]   | llama-8b      |        34.02 ms | 28.60 ms |   5.42 ms

[yuhaohaoyu]

@lasch Thanks for the comments. addressed just now. Below are some details to help quick check

---

Architecture change

Before (3 primitive families, 10 entry points)

_MemAccessor                        HBM/LX Simulators         Free functions
─────────────                       ──────────────────         ──────────────
.read(n, dtype)                  →  .read(...)              →  _read_flat()
.write(data)                     →  .write(...)             →  _write_flat()
.gather(offsets, dtype)          →  .gather(...)            →  _gather_from()
.scatter(offsets, data, dtype)   →  .scatter(...)           →  _scatter_into()

After (1 primitive family, 6 entry points)

_MemAccessor                        HBM/LX Simulators         Free functions
─────────────                       ──────────────────         ──────────────
.read(n, dtype, offsets=None)    →  .read(..., offsets=None) → _read_flat(..., offsets=None)
.write(data, offsets=None)       →  .write(..., offsets=None)→ _write_flat(..., offsets=None)

• offsets=None → contiguous (existing behavior, unchanged hot path)
• offsets=array → sparse gather/scatter (same ravel()[offsets] call)

---

Code Changes

ktir_cpu/memory.py

Action	Detail
Modified	_read_flat — added *, offsets=None kwarg; sparse branch does OOB→zero-fill
Modified	_write_flat — added *, offsets=None kwarg; sparse branch does OOB→drop
Deleted	_gather_from (26 lines)
Deleted	_scatter_into (24 lines)
Modified	HBMSimulator.read / .write — added offsets= forwarding
Deleted	HBMSimulator.gather / .scatter (22 lines)
Modified	LXScratchpad.read / .write — added offsets= forwarding
Deleted	LXScratchpad.gather / .scatter (20 lines)

ktir_cpu/ops/memory_ops.py

Action	Detail
Modified	_MemAccessor.read / .write — added offsets= forwarding
Deleted	_MemAccessor.gather / .scatter (18 lines)
Modified	MemoryOps.load branch 1 (offsets) — mgr.gather(...) → mgr.read(..., offsets=)
Collapsed	MemoryOps.load branch 3 (coords) — removed span-read + fancy-index; now linearizes coords then calls mgr.read(offsets=) directly
Modified	MemoryOps.store branch 1 (offsets) — mgr.scatter(...) → mgr.write(..., offsets=)
Simplified	MemoryOps.store branch 3 (coords) — removed read-modify-write; now calls mgr.write(offsets=) directly
Modified	distributed_load — mgr.gather(...) → mgr.read(..., offsets=)
Modified	distributed_store — mgr.scatter(...) → mgr.write(..., offsets=)

---

Performance Verification

$ uv run pytest tests/ -v
========== 1336 passed, 10 skipped, 13 xfailed, 22 warnings in 33.51s ==========

Benchmark (no regression — within run-to-run variance):

Workload         | T_indirect_load (before) | T_indirect_load (after)
-----------------+--------------------------+------------------------
moe-262K         |                  2.45 ms |                 2.40 ms
moe-1M           |                 10.36 ms |                10.18 ms
paged-attn-256K  |                 34.61 ms |                33.24 ms
sparse-attn-32K  |                  0.19 ms |                 0.17 ms
sparse-attn-128K |                  0.29 ms |                 0.26 ms
retro-attn-16K   |                  0.48 ms |                 0.41 ms
multi-head-4M    |                 35.70 ms |                34.02 ms

---

Overall Change set

• −143 lines net across 3 files (140 insertions, 283 deletions)
• −4 API entry points per layer (10 → 6)
• Branch 3 in MemoryOps.load eliminated (coords and offsets paths converge)
• Store coord path simplified (no read-modify-write needed)
• OOB semantics unified in one place (_read_flat / _write_flat)

---

Comment and test cleanup

Stale references to the deleted gather/scatter/_gather_from/_scatter_into
API were updated in tests/test_blocked_indirect_fast_path.py:

Old	New
class TestScatter	class TestSparseWrite
"""scatter() writes only the targeted offsets..."""	"""write(offsets=) writes only the targeted offsets..."""
test_hbm_scatter_sparse	test_hbm_write_sparse
test_lx_scatter_sparse	test_lx_write_sparse
test_scatter_gather_roundtrip	test_write_read_sparse_roundtrip
class TestGatherScatterOOB	class TestSparseOOB
"""Verify that gather zero-pads OOB offsets and scatter drops them."""	"""Verify that read(offsets=) zero-pads OOB and write(offsets=) drops them."""
test_gather_oob_returns_zero	test_read_oob_returns_zero
test_gather_all_inbounds	test_read_all_inbounds
test_scatter_oob_dropped	test_write_oob_dropped
test_scatter_all_inbounds	test_write_all_inbounds
test_gather_primitive_wraps_negative	test_sparse_read_wraps_negative
# _gather_from wraps silently (section comment)	# _read_flat(offsets=) wraps silently
# gather returns zero ... scatter silently drops (section comment)	# read(offsets=) returns zero ... write(offsets=) silently drops
# Scatter primitives: ... (section header)	# Sparse write primitives: ...

No tests removed — all test real behavior that still exists under the new API.

---

---

shared-view bug fix

The iterator lookup used sub["index_view_idx"] as key into idx_iters, but
idx_values is keyed by subscription index (sub_i). When two subscripts share
the same index_view_idx, both would consume from the same iterator, exhausting
it at ⌈K/2⌉. Fixed with a sequential indirect_counter:

indirect_counter = 0
for sub in iat.dim_subscripts:
    if sub["kind"] == "indirect":
        raw_idx = int(next(idx_iters[indirect_counter]))
        ...
        indirect_counter += 1

[lasch]

Thanks for closing this maintainability gap. 👍

One observation was made: 'No test happens to exercise the fast path. The fast path, as currently gated, is unreachable by every existing indirect-access example in the repo, including the paged-attention/paged-tensor ones that read like its intended use case" (all examples fail the has_direct_expr gate).

I'd say, we can do this separately.

There's another maintainability gap raised about the 2 separate code paths that still are separate.

M1 — unchanged, not resolved. The general path (_resolve_idx_reads→_build_indirect_coords→_flat_memory_offsets) and fast path (_prepare_dep_var_sub_space→_runtime_read_and_expand_sub_space→_gen_offsets_vso_space_via_broadcast) remain two independent pipelines up to offset computation; this commit only touched the terminal read/write step downstream of both, and — per the addendum's own analysis — neither of fabianlim's proposals touches this pipeline anyway.

I'd say here too, should be done separately, because:

A rough high-level outline of what closing M1 (the two independently-maintained offset-computation pipelines) would take, grounded in the current code in Source/ktir-cpu/ktir_cpu/ops/memory_ops.py. M1 is assessed as generally fixable, but only via a substantial refactor of both the PR #188 additions and pre-existing code.

I'm attaching the details as a file:
pr_188_recommendations.md

[yuhaohaoyu]

Thanks @lasch for review comments. I agree to treat the 2 issues as independent and deal with them separately.

I am thinking of open another issue to track these, and close the existing feature-dev-driving one.

[lasch]

Yes. Please open new issues for the 2 work items.

1. tests/examples to exercise the fast path
2. real combination of the fast path into existing code as outlined above

[yuhaohaoyu]

@lasch

Opened issue #212 to track the 2 notes on essential maintainability. Thanks.

[yuhaohaoyu]

@lasch conflict resolved. Please give it a check. Turned out not affected with the HBM --> global renaming, etc.

Conflict resolution

• PR -188 does not touch tests/test_indirect_access.py while it got heavily refactored, accepted the upstream/main one, the rest rebased successfully.
• Run all tests
• Run benchmark compare, no regression

Benchmark compare, no regression

bench_indirect_load_path_compare — indirect-load-path-timing
=================================================================================================================================================
Workload         | Access expr               | Source shape         | Result shape       | Model         | T_indirect_load |   T_load | T_offsets
-----------------+---------------------------+----------------------+--------------------+---------------+-----------------+----------+----------
moe-262K         | X[IDX[e], m, n]           | [128, 256, 128]      | [8, 256, 128]      | deepseek-v2   |         2.54 ms |  1.96 ms |   0.58 ms
moe-1M           | X[IDX[e], m, n]           | [128, 1024, 128]     | [8, 1024, 128]     | deepseek-v2   |        10.07 ms |  7.95 ms |   2.13 ms
paged-attn-256K  | cache[BT[d0], d1, d2, d3] | [1024, 16, 128, 128] | [16, 16, 128, 128] | llama-8b      |        34.88 ms | 29.57 ms |   5.31 ms
sparse-attn-32K  | X[P[e], T[e], h]          | [1024, 512, 64]      | [4, 16, 64]        | bigbird-base  |         0.20 ms |  0.04 ms |   0.16 ms
sparse-attn-128K | X[P[e], T[e], h]          | [4096, 512, 256]     | [4, 16, 256]       | bigbird-large |         0.26 ms |  0.11 ms |   0.16 ms
retro-attn-16K   | mem[doc[d], chunk[c], h]  | [1024, 128, 128]     | [32, 4, 128]       | retro-7b      |         0.55 ms |  0.13 ms |   0.42 ms
multi-head-4M    | X[E[e], H[h], m, n]       | [128, 64, 256, 256]  | [8, 8, 256, 256]   | llama-8b      |        34.93 ms | 29.68 ms |   5.25 ms

[lasch]

/lgtm