[yangzhuxinyzx]

Purpose

Flash-Next TP4 on 4×V100 16 GB retained duplicate workspaces and full HC/PLE temporaries, leaving too little space for long prefill. Bound ownership and lifetimes while preserving 80% GPU / 20% host experts, 512 MiB allocator reserve, host main KV with GPU compressed index history, 256K total context, prefix caching, tools/thinking and FULL MTP4.

Reuse the existing QSA, HC, PLE and GGUF families. Qualified HC prefill keeps rows local, including MTP FC and PLE convolution outputs; the target's final HC gather writes its persistent MTP buffer directly and preserves capacity tails. Share identical RoPE, serial QSA top-k and target/draft hidden storage; size BLAS scratch by admitted projections. Preserve decode policy and qualified fallbacks. The versioned launcher enables the profile, including one page-rounded shared stage that covers FP16 draft history through 262144 tokens.

The module memory/speed comparison also exposed a correctness bug: generic GGUF MMQ skipped experts above 255 although the MTP has512. Pass the weight tensor's actual expert count through the existing kernels, widen expert byte-offset multiplication, and retain the public op ABI and quantization.

Test Plan and Results

• Imported Claude's two commits with our DCO sign-off; integration base remains6f3557c973b4602905037e086fd4acba9b0d867f. Measured runtime head81edd142bd1d48e8fa428d24f17b034348d20ade; the subsequent documentation commit records ordinary acceptance. Existing PR1210 owns this scope; no duplicate PR. AI assistance was used; human review remains required.
  -76 focused memory/ownership tests passed. TP4 target and MTP FC/HC adapters passed six query shapes4096/4097/7344/7345/8192/16384, padding, complete outputs/state and changed-input CUDA graph replay under the original paired0.001 numerical bound. Synthetic4096-row MTP adapter peak delta500.18→180.01MiB; this is not whole-model HBM evidence.
• Normal CMake _C_stable_libtorch rebuilt from the owned source on CUDA12.8/Torch2.10.0+cu128/SM70. Installed standard library SHA256699d904f5d058d6f401930d9ab076160dfea69af66838349ffeb817fda488a44; other11 normal artifacts unchanged from6ca4143c747898ddc3f09895439a9c0359e1973e. All12 dependency audits use only standard CUDA/Torch/system libraries, with no private DSO/preload.
  -18 additional expert-bound, full Q8 MoE and staging-budget cases passed: Q8 E128/256/257/512, Q4_0 E257/512, upper-half expert IDs and changed-input/routes graph replay against dequantized FP32 references; full MTP Q8 crosses the64/65-token dispatch boundary. The original benchmark's uninitialized/nonfinite failure is retained.
• Corrected Q8 chunk outputs were bitwise equal, but1024 tokens cost about1.9ms more per7344/8192-row call; retain2048 for throughput.
• Real FP16 draft host_kv_forward, mapped host backing, page816,513 queries: the old byte budget admits only130560 tokens. At130561 tokens, old fallback≈203ms versus enlarged staging≈31ms; at262144 tokens,≈225ms versus≈59ms. Maximum paired row L2<0.000345 and absolute error7.63e-6, within predeclared0.001/0.002 limits. One shared stage budget525504 costs another130.60MiB; hot KV/state blocks/expert fraction remain unchanged.
• Retained pre-expansion diagnostic snapshots at matching240K completion: idle allocated13946.25→13532.68MiB, peak14345.70→13912.90MiB. Those are413.57/432.80MiB reductions, not the sum of component estimates. The final ordinary process confirms13663.285MiB allocated per GPU,282.966MiB below the matching old13946.251MiB baseline.
• Actual Inductor/VLLM gather code generation showed no full-buffer clone: extra allocation0MiB for aligned7344/8192 rows and0.0586MiB for7345 rows. Strict fresh-process AOT source/rebinding qualification passed on all4 ranks; artificial fixture/cache failures are retained.

Remaining Qualification

The one-boot ordinary Worker acceptance passed:52 sequential requests (13 each at actual32774/65604/131124/261044 input tokens), three long-context retrieval/arithmetic/natural-EOS checks through260977 input tokens, and exact262139 input+5 output=262144 total capacity. No OOM, hang or automatic restart. Cold-prefix32000-word trials had actual34606–34688 input tokens: TTFT6.525–6.576s, median6.5537s and5284.2tok/s. External200ms driver samples showed minimum778MiB free on every GPU after readiness. Startup driver free briefly reached18MiB, including allocator reservations; the post-ready gate does not establish an all-startup physical floor. No profiler or memory observer ran during this ordinary acceptance.

The user requested immediate public API access, so the newly started short profiler was cancelled during loading and ordinary serving was restored. The public tunnel passed authenticated model discovery,401 rejection, arithmetic SSE (first reasoning/text1.59s) and tool-call checks; one initial tool transport disconnect was retained and its separate retry passed. This is a public smoke, not a stability or verifier-performance claim. Exact target-GPU-start→latest-TP4-state-commit costs across32/64/128/256K remain pending. No all-context16ms claim is made; HTTP round estimates and old partial traces are excluded.

Raw reports, failed paths and diagnostic snapshots: /home/test/flashnext-256k/codex-abliterated-20261014/decode-parity-20261015/ and /home/test/flashnext-256k/claude-memdiag/; implementation/range/fallback notes in benchmarks/flashnext_prefill/HANDOFF.md.

[hdq66666]

The SM70 PLE prefill fast path in _short_conv_dilated_prefill_batched (as of 36d2077, and #1205 at 31edb1f) is skipped for every step that holds two or more prefills, so those steps take the padded fallback.

Why: in Mamba align mode state_indices_tensor_p is column 0 of a (requests, 1 + speculative blocks) table, a strided view (stride 5 with MTP4) once a step holds more than one prefill. has_initial_states_p can be strided too. The admission requires both to be contiguous:

and state_indices_tensor_p.is_contiguous()
and has_initial_states_p.is_contiguous()

So a pure-prefill step with ≥2 prefills (num_decodes == 0) never reaches dilated_prefill_short_conv. With the channel chunks from #1209 the fallback no longer runs out of memory, but it is still the slow path.

We hit this on 4×V100-SXM2-32GB, TP4, Flash-Next NVFP4 + MTP4, with no chunk bound in the fallback. A long cold prompt plus a second cold prompt in the same step produced an 8160 + 32 token step. The padded fallback allocated num_prefills × max_len × 10240 FP16 buffers several times over, and the engine died on a 320 MiB request.

Fix we run in production: drop the two contiguity checks from the admission, and pass contiguous copies to the kernel (a handful of int32/bool values per step):

state_indices_tensor_p[:num_prefills].contiguous(),
has_initial_states_p[:num_prefills].contiguous(),

The kernel itself is unchanged; our copy is byte-identical to the one in this branch.

Regression test: test_sm70_fast_path_takes_strided_state_slots, with shapes [8160, 32] and [40, 3, 17].

• The slot table is built as a stride-5 column view.
• The test asserts that the fast path is taken and that it is bitwise equal to the padded oracle (output and state).
• It fails before the change and passes after (SM70 bitwise; A100 within tolerance).

With this change and the #1209 channel chunks plus the score-tile dels from 996e3a1, in our deployment:

• the two-prefill steps of a 150k + 35k + 2-decoder run peak about 160 MiB lower on each rank: forward +1279 → +1119 MiB at 8170 tokens × 4 requests;
• single-prefill peaks and decode step times are unchanged.

Happy to open a small PR with the fix and the test if that helps.

(Investigation and write-up AI-assisted; reviewed and tested by me.)