[FelineStateMachine]

What & why

Motivation. On a DGX Spark (GB10, sm_121), Edge0-145B NVFP4 prompts hang in prefill: the serve watchdog fires with no progress for 60 s during a request (reading the prompt (batched): layer 0 of the prompt chunk from token 0) and restarts the engine. It happens on most prompts but not all, so it looks random. The only workaround is STRATA_LATENT_FUSED_DU=0, which drops prefill from ~900 to ~210 tok/s.

Cause. In latent_dense_kernel<true> (the latent D projection, in moe_fused.cu and its copy in moe_fused_q3k.cu), the epilogue does

const int r_t = nb0 + 8 * n + 2 * tig + qc;   // the token's tile row
if (r_t >= nrows) continue;
...
am = fmaxf(am, __shfl_xor_sync(0xffffffffu, am, 16));   // full-warp shuffles

When the last 64-row tile is only partly filled, lanes whose row is past the end skip the full-mask __shfl_xor_sync calls that the other lanes of the warp wait in. That is undefined behaviour; on sm_121 the warp deadlocks. Whether a prompt hits it depends on its token count mod 64, hence the intermittency.

Change. Every lane runs the reductions; only lanes with a real row store (const bool valid = r_t < nrows; and if (!valid) continue; after the shuffles, before the stores). The shuffles (xor 4, 8, 16) only combine lanes with the same tig, so the same token row: a row past the end only ever feeds itself, and the tile's spare rows repeat its last row (srow), so their values are finite. Rows of full tiles are unchanged bit for bit (checked below). Two commits:

1. fix(strata): the fix in both kernels (4 lines each).
2. test(strata): lat_dense_d_tail_test, a regression test registered with CTest.

No existing issue; I can open one if you prefer to track it there.

Area

• Engine, CLI or server (edge0-infer/)
• Model asset tools (edge0-infer/tools/assets/)
• Chat app examples (edge0-infer/examples/)
• Model docs (docs/)
• Other: edge0-strata/ (prefill kernels)

How was it tested?

Environment: NVIDIA DGX Spark (GB10, sm_121, 128 GB unified memory), DGX OS 7.2.3 / Ubuntu 24.04, kernel 6.17 (nvidia), driver 595.91.07, CUDA 13.0 toolkit. Base commit b7900cb. Engine built per edge0-strata/EDGE0.md with -DCMAKE_CUDA_ARCHITECTURES=121.

Check	Before (b7900cb)	After
lat_dense_d_tail_test (both kernels, Q8_0 / Q4_0 / q3k colscale, M = 1..127)	hangs at the first partial tile (M=1), 3 of 3 runs	PASS
compute-sanitizer --tool memcheck on the test	n/a (hangs)	0 errors
compute-sanitizer --tool synccheck on the test	hangs first, reports nothing	0 errors
Full tiles (M=128), old vs new kernel output, all 4 variants		byte-identical (8.3 MB, same SHA-256)
Edge0-145B NVFP4, 10 prompts of 4K-32K tokens, fused D/U on	most hang (watchdog restart)	10 of 10, prefill 850-1030 tok/s
Edge0-145B NVFP4, agent client with ~38K-token prompts, multi-turn	hangs, or ~210 tok/s with STRATA_LATENT_FUSED_DU=0	completes; 38K tokens read in 41 s (940 tok/s)

Not covered: I only have sm_121, so this was not run on RTX cards. There it is the same undefined behaviour that happens to complete; full-tile output is unchanged and the partial-tile rows are the same rows, so no change in results is expected. The q3k kernel is covered by the unit test only (no Q3_K model run end to end).

cuda-gdb on the hung test, unpatched (built with -lineinfo)

[Switching focus to CUDA kernel 0, grid 5, block (0,0,0), thread (0,0,0), device 0, sm 0, warp 2, lane 0]
0x0000000325b0aa40 in strata::prefill::fused::(anonymous namespace)::latent_dense_kernel<true><<<(48,1,1),(512,1,1)>>> ()
  at .../sm_30_intrinsics.hpp:449 in __shfl_xor_sync inlined from moe_fused.cu:1170
  Kernel Parent Dev Grid Status       SMs Mask  GridDim  BlockDim Invocation
*      0      -   0    5 Active 0x000003030303 (48,1,1) (512,1,1) strata::prefill::fused::(anonymous namespace)::latent_dense_kernel<true>()
  BlockIdx To BlockIdx Count   State
   (0,0,0)     (7,0,0)     8 running
  Wp Active Lanes Mask Divergent Lanes Mask          Active PC Kernel BlockIdx First Active ThreadIdx
Device 0 SM 0
   0        0xeeeeeeee           0x11111111 0x0000000325b01240      0  (0,0,0)               (65,0,0)
   1        0xeeeeeeee           0x11111111 0x0000000325b01240      0  (0,0,0)               (97,0,0)
*  2        0xeeeeeeee           0x11111111 0x0000000325b01240      0  (0,0,0)                (1,0,0)
   3        0xeeeeeeee           0x11111111 0x0000000325b01240      0  (0,0,0)               (33,0,0)
   4        0xffffffff           0x00000000 0x0000000325b01470      0  (0,0,0)              (192,0,0)
  ...
  15        0xffffffff           0x00000000 0x0000000325b01470      0  (0,0,0)              (416,0,0)

M = 1: only tig == 0 lanes (0x11111111) have a row; they wait in the shuffle at moe_fused.cu:1170 while the other lanes (0xeeeeeeee) have left. Warps 4-15 (nb0 >= 16, on0 false) wait at the next item's __syncthreads(). The same pattern was captured from the engine itself during a real prompt (block 16, lanes 0x88888888 split from 0x77777777).

Test output

Before:

FAIL fused Q8_0 M=1: the launch did not finish in 10 s (hang)

After (plain, under memcheck and under synccheck alike):

ok   fused Q8_0: M = 1..127, rows identical to the full-tile run, nothing written past row M
ok   fused Q4_0: M = 1..127, rows identical to the full-tile run, nothing written past row M
ok   q3k Q8_0: M = 1..127, rows identical to the full-tile run, nothing written past row M
ok   q3k Q8_0 colscale: M = 1..127, rows identical to the full-tile run, nothing written past row M
PASS
========= ERROR SUMMARY: 0 errors

Engine, Edge0-145B NVFP4 (serve.log)

Before:

strata serve: no progress for 60 s during a request (reading the prompt (batched): layer 1 of the prompt chunk from token 0) - stopping the engine so the server starts it again (issue #29)
strata serve: stall report (engine 0.1.55): stage "reading the prompt (batched): layer 1 of the prompt chunk from token 0" for 60 s; 0 layers served since the last finished step (0 = stopped, more = slow)
  expert pool: epoch 0, batch epoch 0: 0 of 0 jobs claimed, 0 done; 18 of 18 workers parked, 18 sleeping; mode 0

After:

strata serve: prompt 4212 tokens = 0 reused + 4212 read in 4678 ms (900.5 tok/s), 128 generated in 4030 ms (31.8 tok/s)
strata serve: prompt 8471 tokens = 0 reused + 8471 read in 8549 ms (990.8 tok/s), 128 generated in 4447 ms (28.8 tok/s)
strata serve: prompt 16720 tokens = 0 reused + 16720 read in 16354 ms (1022.4 tok/s), 128 generated in 4229 ms (30.3 tok/s)
strata serve: prompt 33450 tokens = 0 reused + 33450 read in 32459 ms (1030.5 tok/s), 128 generated in 4204 ms (30.4 tok/s)
strata serve: prompt 38432 tokens = 0 reused + 38432 read in 40869 ms (940.4 tok/s), 67 generated in 1647 ms (40.7 tok/s)

Checklist

• ctest --preset cpu -L "L0|L1" passes locally in edge0-infer/ (engine changes): not run, no edge0-infer/ changes
• Ran the higher test levels this change affects: on-device runs of the model above
• Added/updated tests for new behavior (lat_dense_d_tail_test)
• Updated relevant docs / READMEs: none affected

AI disclosure

The diagnosis (the cuda-gdb captures), the patch, the regression test and this description were prepared with an AI coding agent (Claude, by Anthropic) at my direction, on my hardware. Every result above comes from runs on that machine; logs are trimmed to the relevant lines. The commits carry an Assisted-by: Claude (Anthropic) trailer.

[shuk777]

Thanks for helping to debug! LGTM