npu: serve GR00T N1.7 on Ascend 910B4 (396.67 → 52.40 ms, cosine 0.9999570) - #200
Merged
Merged
Conversation
LiangSu8899
marked this pull request as draft
September 12, 2026 18:14
A native Ascend path for GR00T N1.7 -- 3 B of Qwen3-VL ViT, a truncated
Cosmos-Reason2 language model, a VL adapter and a 32-layer DiT action head --
routed as ("groot_n17", "torch", "npu"). The whole model is one captured graph
with three inputs, so a frame refills the camera patches, the robot state and
the noise draw and replays.
396.67 ms of eager reference becomes 52.40 ms, 7.6x, at a combined
denormalized-action cosine of 0.9999570 against the original FP32 reference with
the initial noise pinned. Of that, 48.02 ms is the graph and the rest is host.
The 1-D gripper channel reads 0.9956549 and is not a useful gate, being nearly
constant over a trajectory; end-effector and joint targets are 0.9998915 and
0.9999829. Three hundred replays of the captured action chain give one distinct
answer, and the replay is bit-identical to the same graph run eagerly.
torch.compile was not measured on this backend, so the judged baseline is eager;
on CDNA4 the same model compiles to 77.9 ms against eager's 67.9.
Three native units, in csrc/npu/kernels/groot_n17/ because each serves one
model's fixed geometries. Each exists because the vendor operator at those
shapes is dominated by fixed cost rather than by work.
The DiT's attention drives Mmad directly and keeps a whole head in L0 untiled:
41 query rows against 41, 13 or 448 keys over 32 heads of 48 channels, where the
vendor's prompt flash-attention charges 42 to 49 us for every one of those and
the cost barely moves between 13 keys and 448, while the 41x41 case moves 378 KB
and does 10 MFLOP. The row softmax runs over a whole score plane with no
vector-to-scalar round trips. 14.9 us against 42.0, cosine 0.999996. The value
operand is transposed on the way from L1 to L0B rather than by a separate
permute, which is bit-exact and removes two launches a layer.
The fused add-and-normalise is a draw on time against the vendor's fused form
and is here for its arithmetic: the residual sum is rounded to BF16 once,
because that value is what the next block carries forward, and the
normalisation runs in FP32 from it rather than rounding a second time. It also
adds the bias of the projection that produced its branch, and adds it to the
branch and rounds it there, which is where the biased matmul it replaces
rounded -- that placement is worth 0.9999570 against 0.9999531 end to end.
The evaluation image transform's resize reproduces OpenCV's enlarging INTER_AREA
bit for bit: zero differing samples over twenty-eight million, for 0.30 ms of
device against 7.0 ms of host. A torch version of the same arithmetic measures
8.9 ms, because an integer right shift on this part runs at 15 GB/s where a
multiply runs at 151, and the shifts cannot become float multiplies outside a
kernel -- the vertical pass forms a 26-bit product and FP32 has 24 bits.
Shrinking is a different branch inside OpenCV and is refused rather than
approximated.
Everything else runs on torch_npu operators with the weights in Ascend
fractal-NZ layout.
The build selects per model: FLASHRT_ENABLE_NPU_PI05 (ON) and
FLASHRT_ENABLE_NPU_GROOT_N17 (OFF), so the default emits exactly the four
shared objects it did before these units existed. A deselected model's standard
artifacts are removed rather than left behind, because a selection that only
decides what to add leaves a directory disagreeing with the selection that
produced it and a loader would bind the leftovers. At least one model has to be
selected and a switch that is neither ON nor OFF, empty included, is refused by
name. The published ABI stays at 1: adding units with new symbols changes no
contract behind an already-published entry point, and bumping would refuse every
shared object already built.
Every native entry point here is reached with a raw address, so
flash_rt/npu/core/operands.py checks device, dtype, shape, row pitch and
contiguity in front of each one and names the operand that was wrong. Residency
is a separate check because "the same device as that other operand" passes a
launch whose operands are consistently on the host and still faults, and camera
frames are on the host by default.
The checkpoint description moves to flash_rt/models/groot_n17/weight_spec.py,
where the backends that already share it can reach it without one of them
importing another's frontend; the former module re-exports it, so no existing
importer changes.
The tier is BF16. 910/A2 parts have no FP8 tensor hardware, and an INT8 request
is refused rather than downgraded. INT8 on the DiT was built and measured and is
not here: the quantised GEMM wins 3.05 us a call while the kernel types it adds
to the block's per-layer loop make the launches already in that loop slower,
including ones it never touches. On this part a kernel's cost includes what else
is interleaved with it, about 4.5 us a launch.
Tests run on any host except one file: the frontend's routing and refusals, what
a build selection compiles and removes -- driven against a fake toolchain, so the
script carries no test hook of its own -- the operand checks, each unit's
identity check and the refusal an unbuilt one gives, the action decode and state
encode including every representation they refuse, the weight specification, and
the image transform's tap construction against OpenCV itself at four enlarging
geometries. 146 passed and 52 skipped with no Ascend device present;
tests/test_npu_groot_n17_device.py holds the kernels' numerics and needs the
part.
LiangSu8899
force-pushed
the
feat/npu-groot-n17-native
branch
from
September 13, 2026 10:36
4bf088f to
efacaf8
Compare
LiangSu8899
marked this pull request as ready for review
September 13, 2026 12:52
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a native Ascend 910B4 path for GR00T N1.7 (Qwen3-VL ViT + truncated
Cosmos-Reason2 LLM + VL adapter + 32-layer DiT action head), routed as
("groot_n17", "torch", "npu"). The whole model is one captured graph with threeinputs — camera patches, robot state, noise draw — plus three native kernels: the
DiT's attention, a fused add-and-normalise, and the evaluation image transform's
resize.
Results
Full frame, end to end: one observation in, denormalized robot-space actions out.
Ascend 910B4, CANN 8.5.2, PyTorch 2.7.1 + torch_npu 2.7.1.post2.
torch.compilewas not measured on this backend, so the judged baseline is eager.Per modality: end-effector 0.9998915, joint 0.9999829, gripper 0.9956549 (nearly
constant over a trajectory, not a useful gate). 300 replays of the captured chain
give one distinct answer, bit-identical to the same graph run eagerly.
Build
Default is unchanged — the four shared objects Pi0.5 loads, ABI still 1. This
model's units are opt-in:
Documentation
docs/deployment_npu_groot_n17.md— usage, the prompt bundle's keys, what eachkernel does and why, precision, tests
docs/deployment_npu.md— the shared Ascend buildcsrc/npu/README.md— the three entry points, their envelopes and refusalsTests
python -m pytest tests/test_npu_*.py -q→ 146 passed, 52 skipped without adevice.
tests/test_npu_groot_n17_device.pyholds the kernels' numerics and needsthe part.
Draft: the figures above were measured before this branch's last revisions
(shared-object names, operand checks, build selection, the frontend) and are being
re-taken on this head.