[FEATURE] NKI interpreter by latentCall145 · Pull Request #206 · Deep-Learning-Profiling-Tools/tilelens

latentCall145 · 2025-10-29T21:59:44Z

sorry for big diff :(

Adds the NKI frontend to Triton-Viz:

added a NKI interpreter to apply language-level patching (e.g. allow NKI kernels to be executed line-by-line with NumPy; modifies nki.py)
AST rewriting to convert NKI loads/stores to masked loads/stores (modifies nki_extract_slice.py)
Masked load/store implementation (modifies nki_masked_load.py)
added backend (one of "nki" or "triton") as arguments for patching (modifies patch.py, client.py)
added tracer callbacks for masked loads/stores (modifies data.py, tracer.py)
separate Triton and NKI Trace objects (modifies trace.py)
a lot of visualization stuff (credit: @gujialiang123)

Added Tests

nki-examples/: demo programs to try out the visualizer (currently, only nki-examples/matmul.py is supported)
tests/test_nki.py: basic NDArray slicing
tests/test_masked_load.py: make sure masked load/store works correctly

mark14wu · 2025-11-07T02:40:31Z

Hi @mark14wu , can you also review the PR?

Yes. I have dropped some comments, but haven't finished reviewing all files.

mark14wu · 2025-11-07T04:59:04Z

Some tests in CI is not working.

Traceback:
/opt/hostedtoolcache/Python/3.10.19/x64/lib/python3.10/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
tests/test_wrapper.py:8: in <module>
    from triton_viz.wrapper import SANITIZER_COMMAND, PROFILER_COMMAND
triton_viz/__init__.py:1: in <module>
    from .core import trace, clear, config
triton_viz/core/__init__.py:1: in <module>
    from .trace import trace, clear
triton_viz/core/trace.py:7: in <module>
    from ..clients import Sanitizer, Profiler, Tracer
triton_viz/clients/__init__.py:1: in <module>
    from .profiler.profiler import Profiler
triton_viz/clients/profiler/profiler.py:1: in <module>
    from ...core.client import Client
triton_viz/core/client.py:8: in <module>
    from .patch import (
triton_viz/core/patch.py:61: in <module>
    from triton_viz.core.nki import nki_builder
triton_viz/core/nki.py:3: in <module>
    import neuronxcc.nki.language as nl
E   ModuleNotFoundError: No module named 'neuronxcc'

Please fix.

…ols/triton-viz into nki/patch-dsl

Jokeren · 2025-12-23T01:31:27Z

@mark14wu We need to merge this asap to start code clean up and move towards v3.0

mark14wu

+1

mark14wu · 2025-11-04T02:10:17Z


 launches: list[Launch] = []

+T = TypeVar("T")


Where is T used in this file?

@latentCall145

mark14wu · 2025-11-04T02:17:58Z

+            else:
+                raise ValueError(f"Unknown backend: {backend}")
+
+        # Handle NKI functions specifically


if backend == "nki", the function decorator will return at line 177 and line 186 to 193 will never be executed.

@latentCall145

mark14wu · 2025-11-04T02:18:04Z

+            else:
+                raise ValueError(f"Unknown backend: {backend}")
+
+        # Handle NKI functions specifically


if backend == "nki", the function decorator will return at line 177 and line 186 to 193 will never be executed.

mark14wu · 2025-11-04T02:24:40Z

+        def _extract_user_frames() -> list[traceback.FrameSummary]:
+            stack: list[traceback.FrameSummary] = list(traceback.extract_stack())
+            # drop current frames (this function and callers)
+            stack = stack[:-2]


Maybe we should lift the _extract_user_frames function as a utility, put it somewhere like utils.py and reuse it in both sanitizer and tracer.

mark14wu · 2025-11-04T02:42:22Z

+        def _extract_user_frames() -> list[traceback.FrameSummary]:
+            stack: list[traceback.FrameSummary] = list(traceback.extract_stack())
+            # drop current frames (this function and callers)
+            stack = stack[:-2]


What I mean is _get_traceback_info() in sanitizer.py

mark14wu · 2025-11-04T02:46:17Z

-    call_path: list[traceback.StackSummary] = field(init=False, default_factory=list)
+    call_path: list[traceback.FrameSummary] = field(init=False, default_factory=list)

    def __post_init__(self):


Why we extract call path in __post_init__?

mark14wu · 2025-11-07T05:00:00Z

+        def _extract_user_frames() -> list[traceback.FrameSummary]:
+            stack: list[traceback.FrameSummary] = list(traceback.extract_stack())
+            # drop current frames (this function and callers)
+            stack = stack[:-2]


Please at least move those _extract_user_frames into clients/utils.py.

mark14wu · 2025-11-07T18:24:25Z

+        def _extract_user_frames() -> list[traceback.FrameSummary]:
+            stack: list[traceback.FrameSummary] = list(traceback.extract_stack())
+            # drop current frames (this function and callers)
+            stack = stack[:-2]


@latentCall145

mark14wu · 2025-12-23T02:47:04Z

Before merging, let me fix the conflicts.

Jokeren and others added 30 commits July 28, 2025 10:09

Support NKI

e6e0641

Update

c9b2501

Update

6fe57e6

Update

86d1d6f

Update

7aa3d91

Update

4f8f616

Update

9962f53

Update

25c8ceb

Update

cb07831

Update

1810284

Update

e6a6f91

Update

8406f66

Update

b63e74c

Update

241530b

Update

d28a17e

Update

bc0873a

UPdate

4b6827d

Update

62ade2d

Update

d9d7bf7

Update

998a1a5

Update

2a50a04

Update

8e49f9d

Update

45596d3

Update

a54ac94

Update

f1231fe

Update comments

b1f658a

stuff

e7e2dd2

triton-viz should support cpu to allow development on non-trainium node

b3ee8d9

test kernels to check basic tracing

454ac73

more flexible advanced indexing (just make arange slices an arange)

65e861d

latentCall145 and others added 6 commits November 6, 2025 14:49

debloat nki offsets code

675abab

fix matmul visualization error

7c3bec2

Merge branch 'main' into nki/patch-dsl

86e476a

lint

229d4a9

Merge branch 'main' into nki/patch-dsl

5bf74c8

Merge branch 'main' into nki/patch-dsl

096122c

latentCall145 added 8 commits November 13, 2025 10:34

rename nki_masked_load -> masked_load since it's also used for triton

54559f1

make AWS neuron an optional feature of triton viz (and include CI)

f9c72e3

Merge branch 'nki/patch-dsl' of github.com:Deep-Learning-Profiling-To…

3a00512

…ols/triton-viz into nki/patch-dsl

guard nki test import if NKI not installed

0f6664a

update import

b45cd5e

[CI] fix aws neuron ci

6c01d62

python>=3.10 needed (also 3.9 is EOL)

fa814b5

[CI] make AWS tests reuse torch + triton installations from build job

de93383

latentCall145 force-pushed the nki/patch-dsl branch from a6d38b9 to de93383 Compare November 13, 2025 20:13

latentCall145 and others added 7 commits November 13, 2025 16:12

make nki set_grid_dim consistent with triton api

50a0177

Merge branch 'main' into nki/patch-dsl

24c9abe

Merge branch 'main' into nki/patch-dsl

7260fce

extra nki kernels to viz

76f8999

typo

3dc8c02

Merge branch 'main' into nki/patch-dsl

0817087

Merge branch 'main' into nki/patch-dsl

0eb3f22

mark14wu approved these changes Dec 23, 2025

View reviewed changes

Merge branch 'main' into nki/patch-dsl

232e87b

mark14wu merged commit 4b5ff4a into main Dec 23, 2025
1 check passed

mark14wu deleted the nki/patch-dsl branch December 23, 2025 03:07

Conversation

latentCall145 commented Oct 29, 2025

Uh oh!

mark14wu commented Nov 7, 2025

Uh oh!

mark14wu commented Nov 7, 2025

Uh oh!

Jokeren commented Dec 23, 2025

Uh oh!

mark14wu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mark14wu commented Dec 23, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants