Repository navigation
b12x fused MoE: TrellisMX (procedural NPK) source format with an in-tree P8 runtime - #69
Open
brandonmmusic-max wants to merge 254 commits into
Open
brandonmmusic-max wants to merge 254 commits into
brandonmmusic-max wants to merge 254 commits into
Conversation
Normalize A4 terms and warps to neutral defaults when the prefill threshold disables the feature. Validate explicit controls only for enabled hybrid plans, so stale A4-only environment values cannot block ordinary MoE preparation. Validation: five preparation-control tests passed with CUDA hidden; applicable pre-commit hooks passed. Enabled kernel arithmetic is unchanged.
Reconstruct packed scale bytes in shared-memory stages while retaining one resident weight payload and the dense scale path. Prepare hybrid capability with a boolean control and select precision explicitly at bind time. Share prepared-variant and scale-consumer selection with the NVFP4 prefetch query; use a larger prepared variant when an explicit A4 launch needs its scratch capacity. Validation: 165 tests passed in the combined changed suites, with the two-device ownership case separately passing on GPUs 8 and 9. Pre-commit passed. Actual GLM TP2 layer inputs at M128, M512, M1536, and M3072 matched baseline A4 outputs bitwise and passed changed-input graph replay and allocation checks.
Prepare stage and dense A4 scale consumers without duplicating weight payloads or scratch. Keep stage decoding by default and for live row counts at or below the immutable CSF stage limit. Select dense consumption only when the caller guarantees ordered expanded scales. Share readiness and size selection between binding and the NVFP4 scale-consumer query; precision remains explicit per call. Validation: 115 host precision tests and eight GPU correctness cases passed. GPU cases cover reader selection at 1536/1537 rows, dense/stage bitwise parity, changed-input CUDA graph replay, ordered expansion, and allocation stability. Pre-commit passed. Serving performance qualification remains separate.
Mask the final FC1 half tile and stage gate/up CSF half-slabs using their original exception records. Keep the aligned branch and numerical contract unchanged, without padding or duplicating weight payloads. Bump the compiler identity and tuning query schema. Validate dense and CSF readers with one- and two-plane quantization, changed-input graph replay, payload aliasing, and byte-exact half-slab reconstruction. The float64 oracle also covers synthetic H2560/I320/top-k10 rank geometry.
Remove import-time PTX assembler discovery and environment mutation so Triton planning and child-process compilation retain the same compiler identity. Document explicit Blackwell assembler selection before process startup. Validation: CUDA_VISIBLE_DEVICES=1 /home/luke/projects/vllm-jj-rebase/.venv/bin/python -m pytest tests/test_env.py -q (8 passed); pre-commit run --files docs/installation.rst flashinfer/triton/__init__.py tests/test_env.py (passed). Assisted-by: OpenAI Codex Signed-off-by: Luke Alonso <lalonso@gmail.com>
Expose b12x plans for cache writes with separate key/value dimensions and raw-input depthwise convolution with caller-owned history and scratch. Share convolution kernels with PLE and retain graph-safe dynamic metadata.
Carry an optional fixed attention partition count through QSA planning and execution. Round raw-ring capacity to a page divisor that covers the rollback horizon, including seven drafts on 64-token pages. Add capture and rollback regressions.
Propagate deterministic output requirements through workspace planning and tuning legality. Invoke compiled Triton router launchers using their full positional signature.
Include recurrence policy and intermediate precision in prepared query identity. Keep BF16 chunk execution as the default; provide an explicit FP32 device recurrence for numerical diagnostics and a sequential single-key-split policy for callers requiring fixed reductions.
Retain prepared attention with selectable intermediate arithmetic for identical-input comparisons. Preserve the production default and isolate the eager-BF16 diagnostic implementation in preparation identity.
Split normalization sums and probability writes across CTAs while preserving each logical reduction lane and the 1024-thread final reduction. Reuse histogram storage after its final radix pass.
curand_uniform can return exactly one. Accept positive target-to-draft ratios of at least one in both chain acceptance loops while retaining RNG draws and the strict lower-ratio comparison. Regression cases cover equal distributions, point masses, independent acceptance counts, lower ratios, scalar/device RNG, and graph replay. Targeted GPU qualification: 72 passed.
…e norm Allow exact in-place residual updates while rejecting unsafe output aliases. Specialize vectorized affine loads by gain dtype and centering, preserving BF16 residual rounding. Validation: 36 prepared norm and compilation tests plus the direct CuTe live-count reuse test passed on SM120; changed-input graph replay covers capacities through 4096. Ruff checks pass.
Carry the input row stride through convolution and speculative-history updates so packed QKV/Z outputs require no intermediate copy. Preserve contiguous columns, disjoint mutable buffers, and exact declared strides. Validation: 26 short-convolution tests passed on SM120, including changed-input graph replay in decode, prefill, and mixed modes with 6144 columns of padding. Ruff checks pass.
Make the prepared BF16 and FP32 paths opaque to torch.compile while declaring scratch, output, and recurrent-state writes. Rebind compiler-supplied tensors to the selected plan without changing kernels or precision. Validation: four fullgraph compilation and CUDA replay cases pass across both precisions and 17/65-token inputs, including changed operands, empty requests, checkpoints, exact state/output equality, and allocation-stable replay.
Declare selected logits, IDs, and route weights as custom-op mutations while retaining prepared scoring, bias, and modality contracts. Validate caller options against the selected invocation before dispatch. Validation: all 13 routing tests pass, including fullgraph compilation and changed-input CUDA replay for both score functions, strided outputs, modality controls, and allocation stability.
Register precise output and barrier mutations for expert-presence masks, NVFP4-CSF expansion, and deterministic route reduction. Retain supplied scale-decoder programs and resolve the reduction stream at invocation. Validation: 131 GPU tests passed; one two-device validation test skipped on the single assigned device. Includes compiled captures, changed routes, poisoned outputs, and allocation-stable replay.
Register output-mutating operators around softmax, top-k masking, top-p renormalization, and probability sampling. Preserve kernel selection, RNG tensors, and public return values. Validation: 196 sampling pipeline and RNG tests passed on SM120, including full-graph compilation, CUDA replay, changed parameters, and signed scalar and uint64 tensor seeds.
Expose the prepared writer through a mutating Torch operator, preserving runtime tensor binding and validation. BF16 and FP8 strided cache writes remain visible to ATen tracing. Validation: 13 SM120 writer tests passed, including compiled graph replay with changed slot mappings and strided packed K/V buffers.
Declare the output tokens and acceptance counters as mutable outputs of the existing native verifier. Preserve scalar and device RNG inputs, public signatures, and kernel arithmetic. Sampling compilation, capture/replay, and RNG coverage: 208 CUDA tests pass.
Register the oneshot and fused launcher memo tables with program_cache so deferred compilation closures are discarded when preparation resolves their kernel programs. Validation: PCIe preparation tests (28 passed, 9 CUDA cases skipped); Flare compiled TP2 and TP4 native graph replay passes for a 4096-element BF16 reduction.
Expose scratch, cache, and pooled-selection mutations through custom operations. Reconstruct output views from byte scratch within the trace, and supply planned LSE storage to prefill even when return_lse is false. Preserve decode and extend return conventions. Validation: 8 compiled decode/extend cases with 16/32 heads and up to 4096 rows match underlying prepared launchers exactly with no eager or replay allocator changes; 8 pooled-selection cases pass, including compiled live-metadata replay.
Retain the existing vLLM Triton algorithms behind prepared custom operators with caller-owned cache and sequence-tail storage. Support ordered decode and parallel chunked prefill without runtime allocation or kernel resolution. Validation: seven component tests pass, including nine consecutive 65-token chunks across physical pages. Two 4096-token strict PyTorch byte comparisons remain failing at two one-byte differences; the same inputs match the original vLLM kernel exactly. All seven copied kernel functions retain identical ASTs. Ruff passes.
…h reuse Expose tensor dependencies and mutations through a custom operator. Reset fused merge counters before prepared execution, retain unused scores in planned scratch, and alias constexpr-dead prefill outputs instead of allocating temporary tensors. Validation: six compiled graph cases pass for 32-head decode and prefill through 4096 rows, with changing inputs, lengths, poisoned shared scratch and stable eager/replay allocation counters. Selection comparisons require exact IDs except equal-score cutoff ties and exact scores. Ten existing API and contiguous-scorer checks pass.
Tile the row dimension within one launch, mask partial row tiles, and admit 32-head batches plus the 256-to-512 absorbed query geometry. Preserve packed rowwise weight and scale views, zero-workspace execution, and the fused RoPE query projection contract. Validation: 38 BMM and fused-query tests pass. GLM cases cover TP2/TP4 head partitions, strided token-major storage, partial row tiles and 4096-row compiled graph replay. Numerical error stays within the existing cuBLAS-versus-FP64 bound, guard rows remain intact, and replay allocation counters remain stable.
…ree P8 runtime Add a `trellismx` source format to the experimental b12x fused MoE backend and the P8 runtime that serves it. A TrellisMX expert sidecar stores routed experts as procedural NPK trellis streams (K2 to K4) that decode to E4M3 inside the MMA, with one UE8M0 scale per 32 weights and a coupled H512/H128 boundary built from fixed suh/svh vectors and sign draws. b12x's existing trellis support is the EXL3 path (W4A16) and cannot decode this format. `fused_moe.trellismx.prepare_trellismx_experts` validates a sidecar, failing closed on anything else, and lowers it onto the dynamic W4A8 trellis kernel, which gains a procedural NPK decode arm and a fused input boundary. The small-M arm runs the W4A16 trellis kernels behind the same boundary for up to `smallm_max_tokens` tokens, a plan parameter (default 16; 0 serves every token count on the grouped arm). `fused_moe.trellismx_p8.attach_trellismx_p8_runtime` serves a prepared layer through a P8 runtime whose source files the caller pins by SHA-256; the in-tree runtime is included. Compile workers return every program the runtime launches for calls of 1..capacity rows, keyed by arm; materialization installs them before freeze, and a launch whose program was not declared fails closed. Both paths take their scratch and their output from the caller: the plans declare the bytes, bind carves them, and no call allocates. The P8 launch configuration is a candidate of the component's tuning contract, raced per capacity. The P8 kernels subclass b12x's shared phase, launch and dynamic kernel modules, which gain default-off selections, and the procedural NPK state law lives once, in `moe/_shared/kernels/npk_state_law.py`, for K2 to K5. `_activation_scale_cols_pad` sizes the activation scale plane of `w4a8_mx`, `w6a8_mx` and `w8a8_mx` at K/32, reclaiming the workspace the FP4 width over-allocated. Validation is in the pull request.
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars.
 |
) Prepare fixed 64-CTA, 512-thread joint launches with one active eight-CTA cluster for phased BF16 transfers. Preserve ordinary profile identities and IPC layouts while validating cluster eligibility and residency. Validate profile admission and owner intervals on CPU, and exact CUDA graph replay, startup selection, aliasing, allocation stability, and profile transitions on TP4 SM120. AI-assisted implementation and review.
…ashinfer-ai#43) The paged TMA consumer masks scores beyond cache_seqlens, but the PV tensor-core multiplication still consumes the uncommitted V fragment elements of the last tile. Zero probability does not suppress NaN under floating-point multiplication, so NaN bit patterns in unused V-cache bytes (stale pages, aliased storage) poison live decode outputs. Predicate only the out-of-range B-fragment halves by their K coordinate before the existing BF16 MMA. Live values, including live NaNs, are unchanged; KV storage is not modified. TMA synchronization, mode validation, route selection and the native-FP8 K32 path are unchanged. Reproduction (RTX PRO 4500, SM120; 24Q/4KV, D256, 128-token pages): poison only rows beyond cache_seqlens and replay a captured graph at lengths 1/33/127/128/129/257. A value-only or K+V tail poison makes live decode outputs NaN for BF16 and FP8 E4M3 KV at B1/B4; a key-only tail stays finite. Downstream, unmodified Flare 5f37bd1 serving the official Qwen3.8-27B NVFP4 checkpoint failed 11/72 native requests with an invalid-distribution status; re-exporting with only this change passes 72/72. Validation: 4 tail regressions fail before and pass after; 19 existing paged-attention component cases pass; 48 long-context poisoned-tail cases (4097/4923/65535 tokens, BF16/FP8, B1/B4) pass. Not tested: SM121, the native-FP8 K32 route. Authored by an AI coding agent (Codex, GPT-6) and validated on hardware; committed unmodified by Claude Code, which did not author or review the kernel change.
Keep compiled metadata, bound scratch updates and external schedules on the same dense full-context split policy. Preserve captured capacity and long-context tails. Invalidate costs measured before schedule restoration and mapped-page tuning fixtures. Cover internal and external captured replay at two and eight requests with eight query rows each, heterogeneous lengths, inactive rows and allocation stability. SM121 guards and paired full-round measurements validate the production delta; eight-request integration coverage awaits hardware validation.
Gather historical and current K/V with the existing BF16 intermediate rounding, clamp and FP8 conversion semantics. Preserve negative-row normalization, invalid-row traps and inactive output padding. Write both final staging buffers in one prepared kernel. Exact-byte, alias, mutation and model guards are supplied by the Flare component integration.
Keep scalar calibration and block-FP8 scaling while omitting unit-scale TMA descriptors, transfers and shared storage. Preserve the cp.async path and swapped output coordinates, including existing wide indexing. Refresh affected tuning costs and cover mutable graph replay with zero dummy scales, swapped output tails and up to 64 rows. Integration GPU validation is pending.
Declare gate-first BF16 source rows through the existing blockscaled query. Preserve BF16 SiLU and product rounding, scalar calibration and padded-scale writes. Normal preparation compares complete A16, materialized quantized and fused quantized projections with owned scratch. Cover mutable captured replay, independent materialized references and source validation through 64 rows; integration GPU validation is pending.
Use one lossless packed weight layout across row counts to execute dense NVFP4 projections without activation quantization. Prepare launch geometry, scratch, and compiled kernels before graph capture. Include row-dependent decoding in binary cache keys and exclude unsupported register geometries from tuning candidates. Validation on SM121: seven packed projection CUDA cases passed, covering rows 1, 8, 16, 32, and 64, capacity 2048, mutable inputs, zero rows, and allocation-stable replay. Flare QAD projection tests passed. Candidate enumeration passed for both Qwen dense MLP weight shapes.
Remove the dense packed weight type, conversion, preparation dispatch, and dense-only W4A16 specializations. Dense projections use the canonical NVFP4 representation shared by prefill and decode. Existing MoE W4A16 support remains available. The staged tree restores the implementation preceding the packed dense feature. Import and preparation tests are pending; no GPU validation has run for this removal.
Retain row-fit, head-split and shared-memory limits for larger query capacities. Twelve GB10 captured-replay cases passed at widths 8, 13 and 16, covering causal and noncausal attention, split and unsplit execution, changing pages and lengths, poisoned storage and allocation stability.
Avoid allocating pipeline stages beyond the longest split partition. The GB10 output projection uses 46,080 rather than 58,368 bytes of shared memory per CTA with unchanged registers. Exact-shape and uneven-partition replay checks passed; paired complete decode and prefill measurements showed no regression.
Admit unswapped single-tile native MXFP8 split-K plans through M16 while retaining the existing block-FP8, PDL and fused-activation bounds. Cover mutable captured inputs, row tails, poisoned workspaces and allocation stability across eligible M tiles and both reduction modes. Advance the dense candidate contract to 19 and the combined blockscaled contract to 36 alongside weight-only staging. GB10 full graph pairs improve W13 111.6195 to 110.3203 ms and W16 113.4495 to 111.7010 ms. Prefill measured a 0.22–0.50% cost; 32 split-K replay cases passed on the research source.
Use compact NVFP4 scale staging, narrow K1024 tiles, and persistent CTA budgets bounded by shared storage. Add MXFP8 N32/K256 tuning choices and 256-column paged attention with head-slice-aware split budgets. Keep split-K output storage separate from input batch scheduling and store each complete accumulator once. Preserve FP32 reduction workspace sizing, align direct N32 validation with preparation, and isolate tuning measurements from object prefetch. Include measurement settings in cached selections. Validation on SM121: 256 native/A16 reference and captured-replay cases, 144 ordered preparation cases, and eight direct N32 cases pass. The matching Qwen integration passes compiled W8/W16 and chunked-prefill checks on four Sparks, plus package export and two native serving tests. SM120 and SM103 were not rerun.
Let ActivationSpec retain calibrated A4 prefill on W4A16 packed weights without environment controls. Bind each execution explicitly to A4 or A16 and reject missing calibration or unsupported A4 geometry. Validation: 172 CPU tests pass and five GPU cases are skipped in the CPU run. Twelve unchanged dense-GEMM tuning assertions fail identically on baseline because they omit split_k=16 candidates. Previous SM121 validation covered four GPU cases and the Qwen Flash package with forced A16 decode.
* Apply the NVFP4 four-warp SFB reshape once in dynamic MoE 42ba2f1 made the shared dense _partition_fragment_SFB split the N and K modes for NVFP4 fragments with atom shape (1, 4, 1). The dynamic MoE kernel calls that helper and then performs the same split in its direct_sfb and fc1_direct_sfb branches. For the M16 four-warp NVFP4 candidate the second split unpacks an integer and compilation fails with "TypeError: Value after * must be an iterable, not int", which aborts fused-MoE tuning. Split the modes in the MoE kernel only while the fragment is still rank 2. Add a live-graph oracle test for the M16 four-warp NVFP4 dynamic candidate, which no existing runtime test compiled. * Cover N=160 and varied block scales in the four-warp MoE test Parametrize the M16 four-warp NVFP4 live-graph oracle test over intermediate_size 128 and 160, and over constant and varied E4M3 block scales. The varied scales differ across expert, row, and K block, so a permuted scale layout changes the output. Assert that the prepared plan selected dynamic_mma_warps=4. N=160 prepares separate W13 halves, which disables swapped FC1. The swapped FC1 fragment with the four-warp atom is reached only with CSF inline scales and is covered by the existing 16-4 cases of test_inline_scales_bypass_expansion_under_frozen_graph_replay. --------- Co-authored-by: mark-yong <mark-yong@users.noreply.github.com>
* Share caller-owned storage for PCIe startup tuning Retain fixture arenas, reference banks, and cold eviction inputs through an explicit preparation owner. Omitted storage shares a runtime-owned instance, preserving one eviction allocation per runtime; explicit owners also share across runtimes. Freeze captured addresses and allocate private fixtures for later growth. Validate external eviction storage without copies or fallback allocation. Cover aliased reductions, gather collisions, repeated captured preparation, stable replay allocation, and storage release on TP4 SM120. AI-assisted implementation and review. * Account for startup fixture storage retained by its owner Include shared arenas and reference banks when a producer uses private fixtures after capture. Report owner-held internal eviction storage for warm producers, retain external eviction exclusion, and deduplicate by physical allocation. CPU tests cover retained growth, internal and external eviction, shared allocation keys, and unchanged aggregate accounting before capture. Binding, kernels, capture, and replay are unchanged. AI-assisted implementation and review. * Add reproducible TP4 startup memory measurements Measure retained startup memory with identical message sizes, prepared geometry, resident weights and validated CUDA graph replay. Record baseline and shared-storage results on SM120 without making inference latency claims.
Make logical KV extents mandatory at paged PV helpers and clear invalid sparse MLA payloads and scales before arithmetic. Mask FP8 TMA operands before MMA, including preloaded tiles, and avoid negative-tile loads for empty causal CTAs. Cache storage remains untouched. Preserve live NaNs through FP8 widening, softmax maxima, normalization, and split merging while keeping empty partials neutral. Correct native K32 operand packing and retain accumulated generic-extend state when a partial tile contributes no keys. Add captured poisoned-storage tests with clean arithmetic references, logical-length transitions, live-NaN controls, and cache immutability checks. Repair GLM MLA production fixtures to use writer/reader plans. Validation: installed precommit checks pass, including mypy and Ruff. GB10 SM121 runs passed 88 backend regression cases and 55 Flare component cases. An exported nvidia/Qwen3.8-27B-NVFP4 package served 143 paired 16k/32k prompts (286 HTTP requests) with a 50 GiB budget, 409 evictions, 285 restores, and no request failures; the unpatched package failed at pair 22 under the same workload. SM103 and SM120-only ragged/public-dispatch paths were not exercised. An existing eager-BF16 D72 varlen LSE check fails identically with the unpatched backend: one of 2064 elements differs by 9.97e-5. Its numerical bounds remain unchanged.
* Mask KV below the window start in paged PV paths The per-pair PV mask makes logical extents mandatory but has no lower bound: windowed requests read V below the earliest query row's window start. The score-stage mask already excludes those keys, but a zero probability does not cancel non-finite V in P·V, so non-finite values below the window start propagate to the output as NaN. Add a per-warp window lower bound shared by the warp's PV fragment: rows are packed in token order, so the earliest row gives the smallest window start. Extend and non-windowed callers pass Int32(0). Add a poison regression test for the window boundary, in direct and CUDA-graph-replay modes: storage below the window start must not change any output bit, and a V-only live control at the first windowed key (K finite) must stay visible. Validation on SM120 (RTX PRO 6000); SM121 and SM103 were not exercised: the regression test fails on the unpatched base and passes with the fix in both modes; poison battery 72/72 with the fix and 48/72 without it (all and only windowed modes); 840/840 bit-exact pairs against the unpatched base with the same window settings on clean inputs. * Isolate the V path in the window test's live control; add a CUDA-graph replay parametrization The live-boundary control poisoned K and V together at the first windowed key, so a K NaN reaching the output through the score path could satisfy it without exercising the V mask. Keep K finite and poison V alone. Each parametrization now also runs captured in a CUDA graph and replayed: poisoning between replays must behave exactly as in direct mode, since the graph reads the live cache buffers. * Add aligned-boundary and warp-straddle cases to the window test The straddle cases poison one inter-warp key band with the window start placed at the first warp-aligned token boundary, derived from the prepared plan's head_splits, and skip geometries with no boundary within the width. The aligned cases place the window start on 16-key and page/stage boundaries where an off-by-one in the bound would be invisible at unaligned offsets. --------- Co-authored-by: mark-yong <mark-yong@users.noreply.github.com>
* Keep paged MSA extend loads within the visible extent The V producer passed the packed extent as its valid-row count, so V rows past cache_seqlens in the last selected page were loaded real and reached P*V: scores for those rows are masked by selection, but a zero probability does not cancel non-finite V (NVIDIA/spark-local-inference flashinfer-ai#203). At page size 64 a union block's second subpage can lie entirely past the request; the kernel read its page id one column past the request's page-table row and copied K and V from the resulting wild address, faulting on FP8 and faulting order-dependently in BF16 (NVIDIA/hyperflare-inference#205). Add _msa_v_stage_valid_rows, mirroring the consumer's visible-extent trim (msa_tile_max_visible_len for union tiles, msa_visible_len for per-token selections), and pass it as valid_rows to every K and V paged-tile copy so rows at or past the visible extent are zero-filled instead of loaded. Guard the page-table load on row validity under entries_per_block > 1 so a dead subpage loads no page id and no K/V; at page size 128 the guard is constexpr-folded and the V trim still applies there, so the page-128 specialization differs from base by the trim. Validation on SM120 (RTX PRO 6000), on these exact bytes: initcheck 0 uninitialized reads on both repros (8,192 before); 8/8 crash probes (BF16 page-64 sequence both orders x caching/no-caching, FP8 page-64 repros x both allocator modes); NaN-tail regression 16/16 (8 geometries, direct and CUDA-graph replay, poison between replays) in one process; the new FP8 page-64 dead-subpage test fails on the unpatched base (CUDA_ERROR_ILLEGAL_ADDRESS) and passes with the fix in both modes; suites msa 64 + verify 23 + e2e 2 + unused-KV paths 11, 0 failures. * Cover the FP8 page-64 non-finite tail in the NaN-tail test Add (True, 64, None) to the parametrization: FP8 KV cache at page 64 with the union tile, alongside the existing BF16 page-64 case. The poison loop already handles the FP8 NaN pattern used at page 128. With this entry the test covers BF16 and FP8 at both page sizes across top-k 2/16 and direct plus CUDA-graph-replay modes (20 parametrizations; the full MSA file runs 68 tests). --------- Co-authored-by: mark-yong <mark-yong@users.noreply.github.com>
…r-ai#75) The RoCE runtime exposes the two halves of its pipelined two-shot all-reduce as separate operations. `reduce_scatter` reduces a whole message and keeps this rank's row shard of the sum; `gather_shards` concatenates every rank's shard along dimension 0. Both run over the two-shot path's pipe_send/pipe_recv buffers, chunk flags, counters and grid, so they interleave with one-shot and two-shot reductions inside one CUDA graph. A reduce-scatter followed by a shard gather moves the bytes of one two-shot all-reduce and reproduces its bits, since each shard is reduced by its owner with FP32 accumulation in rank order, which lets a caller run row-local work on its shard between the halves. `should_reduce_scatter` and `should_gather_shards` admit operands whose rows hold whole 16-byte packs, whose (gathered) row count divides by the world size, and whose whole message fits pipe_bytes; like `should_allreduce` they depend on dtype, shape, contiguity and size only. The preparation surfaces `AllReduce.reduce_scatter` and `AllReduce.gather_shards` compile the kernels with the two-shot launchers; `_PreparedRoce` carries them and primes them through `prepared_reduce_scatter_call` and `prepared_gather_shards_call`. Proxy ABI 11 adds control-record op kinds 3 (reduce-scatter: post the phase-0 chunks only) and 4 (shard gather: post the phase-1 chunks only) to `run_pipe_op` through a phase mask and changes no region layout. Runtimes built from this proxy serve packages recording ABI 9, 10 or 11. Validation: the new kernels compile for two and four ranks in BF16 and FP32 on an RTX PRO 6000 (SM120); the proxy compiles with gcc and libibverbs; the CPU RoCE proxy, preparation-compile and route tests pass (test_declaration_rejects_nonserializable_dtypes fails before and after this change). The multi-node test tests/experimental/b12x/comm/test_roce_scatter_gather_gpu.py, which compares both operations with NCCL and with the two-shot all-reduce's bits and replays a graph interleaving every op kind, has not been run: it needs two or more DGX Sparks on a RoCE fabric. Co-authored-by: Jason Cook <delusion@gmail.com>
Comments, docstrings and error messages in the TrellisMX P8 runtime and its tests describe present behavior.
) The bounce reader pins its 16 MiB slot buffer twice. io_uring buffer registration takes a long-term (FOLL_LONGTERM) pin: the kernel faults missing pages in outside CMA pageblocks and ZONE_MOVABLE, and migrates present pages out of them before pinning. cudaHostRegister takes the NVIDIA driver's pin, which is long-term only on x86. On aarch64, pages that the CUDA pin faults in can come from CMA pageblocks, and the CUDA pin keeps them from migrating, so io_uring_register_buffers failed with EFAULT whenever CUDA registered the slots first and one of their pages was in a CMA pageblock. Later rings that reused the same heap range failed the same way. Register the slots with io_uring first and with CUDA second. Teardown order and the copy path are unchanged. Validation: aarch64 host with a 64 KiB-page 7.0 kernel, seven MIGRATE_CMA pageblocks (3.5 GiB), NVIDIA open driver 615.71.09 and an RTX PRO 6000 Blackwell GPU (SM120). A Flare CUDA weight-loading suite that reads every checkpoint through the bounce reader creates 59 rings per run. Without this change, 3 of 5 runs failed bounce initialization with EFAULT. In the two failing runs traced with the migrate:mm_migrate_pages tracepoint, the kernel reported failed long-term-pin migrations from the failing process; the traced passing run reported none. In one failing ring, five pages failed a single-page io_uring registration with ENOMEM while the CUDA pin was held; cudaHostRegister had faulted all five in. Removing that CUDA registration let the same io_uring registration succeed and migrate exactly those five pages. With this change, 19 runs (1121 rings) had no bounce initialization failure and no failed long-term-pin migration. The b12x loader tests test_direct.py (42 tests) and test_ple_reader.py (6 tests) in tests/experimental/b12x/loader pass.
Merges local-inference-lab/flashinfer main at 057e6d1. The one content conflict, in fused_moe/_preparation.py's plan(), is the controls snapshot: main passes the plan's declared A4 prefill capability to _control_snapshot so an explicit a4_prefill pins one activation plane and eight warps, while the branch resolves the TrellisMX small-M policy and appends trellismx_smallm and trellismx_smallm_limit to the same mapping; the result calls _control_snapshot(experts.plan.activation.a4_prefill) and then extends that mapping for trellismx, so both hold. In the same file main's bind() raises when an explicitly requested A4 launch is missing and the branch's binding carries the TrellisMX arm state, and both edits are kept. planning.py keeps main's ActivationSpec.a4_prefill field, its validation and its preparation rules beside the branch's TrellisMX activation and recipe checks; kernel.py keeps main's rank guard on the NVFP4 SFB reshape; README.md keeps main's A4 text and the branch's TrellisMX section.
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
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.
📌 Description
Adds a
trellismxsource format to the experimental b12x fused MoE backend, and the P8 runtimethat serves it, as one in-tree implementation.
A TrellisMX expert sidecar stores routed experts as procedural NPK trellis streams (K2 to K4) that
decode to E4M3 inside the MMA, with one UE8M0 scale per 32 weights and a coupled H512/H128 boundary
built from fixed suh/svh vectors and sign draws. b12x's trellis support is the W4A16 trellis path
(
b12x_trellisand EXL3) and cannot decode this format. GLM-5.3 Flash TrellisMX, quantized directly fromzai-org/GLM-5.3-Flash-BF16, ships its routed experts this way.fused_moe.trellismx.prepare_trellismx_expertsvalidates a sidecar, failing closed on anythingelse, and lowers it onto the dynamic W4A8 trellis kernel, which gains a procedural NPK decode arm
and a fused input boundary.
fused_moe.trellismx_p8.attach_trellismx_p8_runtimeserves a preparedlayer through a P8 runtime whose source files the caller pins by SHA-256. The in-tree runtime is
included; an external runtime with its own compiler can be attached the same way.
For the in-tree runtime, and any runtime compiled through b12x, preparation owns every program these
paths launch; a runtime with its own compiler declares no b12x compile job and owns its programs.
The in-tree P8 runtime states its compile job as a
P8ProgramSpecand declares every arm andreducer that a call of 1 to its capacity in rows reaches; compile workers return those programs
keyed by arm, and materialization installs them in the live runtime before the session freezes
kernel resolution. A launch whose program was not installed fails instead of compiling.
Neither path allocates during a call. A P8 call takes its whole arena from the caller:
P8NativeTPMoE.scratch_specsreports the bytes a capacity's calls need, the plan declares exactly that muchin
MemoryRequirements.scratch, and the runtime carves every buffer out of it; a missing, short ormisaligned (16-byte) arena raises. The result is caller-owned too (
out=is required), so a call neverreturns storage a later call overwrites. The small-M arm's activation, route and output buffers are declared
by its plan (
smallm_scratch_layout) and carved per binding from the variant's arena, so two bindings of onelayer share only read-only programs; for a GLM-5.3 Flash layer at the default limit on a 188-SM GPU the arm
adds 22,356,768 bytes to the arena.
The P8 launch geometry is a prepared choice. A
P8LaunchConfig(direct and grouped tile M, the directcarrier stride, the direct arm's FC1 pipeline depth and warp count, and the resident-grid cap) is a
candidate of the component's tuning contract, so preparation races the configurations that can change a
capacity's call and fixes the winner; the serving configuration is one candidate and the default when a
call names none, and a
P8LaunchConfigpassed toplan_executionpins the choice.P8RuntimeOptionsholds only behavior settings;
GLM53_SERVING_OPTIONSis the GLM-5.3 Flash serving set.plan_execution(..., smallm_max_tokens=N)declares the largest token count the small-M arm serves: 16 bydefault, or 0 to send every token count to the grouped arm. The value is validated, recorded in the plan's
controls and tuning identity and compiled for, and preparation builds the arm's plan and programs for every
planned token count, so bind and capture compile and allocate nothing. No environment variable selects
either arm. A hand-built
TPMoEScratchCapsdeclares the same limit throughsmallm_tokens(default 0: thegrouped arm serves every token count).
The P8 kernels run on b12x's shared modules. The materialized W4A8 phase kernels and their launch are thin
P8 subclasses of the shared ones, and
P8DynamicKernelBackendsubclassesMoEDynamicKernelBackend, whichgains additive selections whose defaults are the code every existing recipe runs: the P8 owners
(
p8_small_m,p8_fc1_tile_n,p8_scale_sandwich,p8_full_coupled),trellis_scaled,trellis_identity_boundary, procedural NPK split materialization at K3 to K5, the raw routing-id form andthe two-warp producer split. The procedural NPK state law exists once, in
moe/_shared/kernels/npk_state_law.py, for K2 to K5 with its two window-extraction laws.One sizing change reaches other recipes:
_activation_scale_cols_padsizes the activation scale plane ofw4a8_mx,w6a8_mxandw8a8_mxat K/32 columns (one UE8M0 byte per 32 values) instead of the FP4 K/16.Their launches read the plane as a flat buffer, so the FP4 width was safe, but it allocated twice the bytes
they read and halved
_dynamic_rows_padded_limit's row budget for those recipes. This reclaims that space;it is not a correctness fix.
The merged branch includes main at
057e6d190. The TrellisMX call sites passraw_routes,split_fc1_scalesandoutput_layoutat their defaults: TrellisMX runs token output with 1-D FC1 gains.The consumer is Flare's GLM-5.3 Flash TrellisMX recipe, proposed separately. GLM-5.3 Flash also
needs "b12x mHC: keep FP32 mixing weights in the TF32 projection", proposed separately, to run its
mHC prefill on the route b12x's autotuner picks.
🔍 Related Issues
review/trellismx-p8-native-20260907): the branch where the P8 kernels were developed. This PRcarries the serving-build version, which adds the RP2 down-remainder kernel, on b12x's shared modules.
🚀 Pull Request Checklist
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.ruff 0.12.8(the version pinned in.pre-commit-config.yaml) checks and formats every Python filethis PR touches;
--all-fileswas not run.🧪 Tests
unittest, etc.).CPU (CUDA hidden) and GPU results are in the evidence block. No numerical bound was relaxed.
🔬 Experimental Track
This PR is experimental: it adds or changes code under
flashinfer/experimental/. Tracking issue: unavailable because issues are disabled; this PR is the review record.flashinfer/experimental/b12x/.tests/experimental/and were validated on the intended hardware; a runnable example is included.flashinfer/aot.py, and no experimental backend is reachable frombackend="auto"withoutFLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1.Owner: @brandonmmusic-max.
Use case: serving TrellisMX releases with b12x fused MoE on SM120/SM121, starting with GLM-5.3
Flash TrellisMX, whose routed experts are procedural NPK trellis streams.
Reason for the experimental path: b12x itself is experimental, and the format has one released model
so far.
Graduation plan: leave the experimental track with b12x's fused MoE, once the GPU cases run in the
experimental CI lane on SM120.
Reviewer Notes
compiled from them gets a different compile fingerprint: NVFP4,
w4a8_mx,w4a8_nvfp4,w6a8_mx,w8a8_mx,b12x_trellis, the W4A16 kernel (which imports the NPK decoder), the TrellisMX arms and P8.Their default specialization is unchanged, and their GPU suites give the same results before and after
(evidence block).
runtime.install_compiled_programs(compile_p8_programs(runtime.program_spec, max_tokens, ordinal), max_tokens),and passes
out=andscratch=(sized byruntime.scratch_specs) on every call.layout before it touches a pointer; the fused MoE paths accept int32 or int64 ids.
the caller's dtype; preparation declares all three.
attaching are seen only by
verify_trellismx_p8_attachment, which synchronizes; graph replaysbypass every check.
release keeps FP32.
encode_csf_scale_planeis only the lossless scale round trip the testsuse.
Evidence: commands, results and precision notes
Commands and results. The results below were measured at the merged head
7b0d0f284, which mergeslocal-inference-lab/flashinfer main at
057e6d190. CPU runs with CUDA hidden, the data runs with the release andtestdata variables; GPU runs on one RTX PRO 6000 Blackwell Max-Q (SM120). The end-to-end line below was measured at
c37d2d9eeon four RTX PRO 6000 Blackwell GPUs at tensor parallelism 4 and is not rerun at the merged head.CPU:
test_trellismx_realdata.py, because the configuredB12X_TRELLISMX_CAPTURESdirectories are unavailable.Without the data variables: 354 passed, 99 skipped.
test_trellismx_p8_runtime.py: 38 passed, 38 skipped; without data 30 passed, 46 skipped.057e6d190without this PR, in the sameenvironment: 163 failed, 2725 passed, 2425 skipped. The failing ids are identical.
GPU:
test_trellismx_p8_runtime.py: 76 passed, none skipped. This includes the tensor-parallel identity tests (the fourTP4 ranks' runtimes of one released layer declare the same programs and the same tuning-race keys) and the arena
poison tests (a call on an arena filled with 0xFF bytes is bit-identical to one on a cleared arena, at seven shapes).
525 passed and 23 skipped out of 548 collected; the 23 are the capture-oracle cases, which skip for the same missing
capture directories.
test_activation_scale_plane.py,test_token_outer_scale.pyandtest_trellismx_smallm_scratch.py: 170 passed, noneskipped.
w4a8_mx,w6a8_mxandw8a8_mxbytecontainers, EXL3, FC2 scale staging): 441 passed, 2 skipped.
w4a8_mxoracle subset (33 cases) and thew8a8_mxsuite (33 cases) pass with boththe FP4 and the K/32 width;
w6a8_mxhas no end-to-end run path in the tree, and its planned width is pinned by aCPU test.
examples/experimental/b12x_trellismx_moe.py, synthetic: small-M m=8 cosine 0.9999981 (bound 0.99995), groupedm=32 0.9998261 (bound 0.9998).
examples/experimental/b12x_trellismx_moe.py --sidecar <released-rank-sidecar-file> --tokens 1 16 17 64:small-M m=1 and m=16 cosine 0.9999991 and 0.9999990, grouped m=17 and m=64 0.9999101 and 0.9999034.
examples/experimental/b12x_trellismx_p8_runtime.py --layer 3 --rank 0 --tokens 1,2,16,17,128,129,1024: theruntime equals the fused MoE call bit for bit at all seven counts, cosine 0.9999976 to 0.9999999 against the
reference.
PreparationSessionwith compile workers, P8-served layer 3 rank 0, 172 launch candidates raced: RESULTPASS at 12 token counts from 1 to 4096. Every row runs through P8 dispatch and matches the standalone runtime bit for
bit, eager and under graph replay; the direct RP2 rows (m <= 16) reach cosine 0.9999975 to 0.9999978 against the
unquantized reference and the grouped rows 0.9999998 to 0.9999999 against the quantized one. Graph replay takes
78 us at m=1, 322 us at m=16 and 3.03 ms at m=4096.
PreparationSessionwith compile workers, the PR's own arms, 16 candidates raced: small-M rows cosine0.9999990, grouped rows 0.9999033 to 0.9999090 through m=256, and m=4096 (first 64 rows checked) on the session's
uniform 0.9999 bound: five runs of these arms gave 0.9998999, 0.9998997, 0.9999002, 0.9999002 and 0.9998997, so it
prints LOW in three of them (see the precision note).
projection fix from an empty kernel cache (all 1,287 preparation requests, every rank declaring the same tuning
races, 122 serving graphs) and served it at TP4 on four RTX PRO 6000 Blackwell GPUs: 443 of 469 reliability
requests at temperature 1 answered correctly, with no request or device errors; decode ran at 189.2 tok/s for one
stream and 751.9 tok/s for 32 streams at context 0 (llm-inference-bench v0.4.29). That run's FlashInfer runtime, apart
from the separately named mHC fix, differed from this branch in comments, one validation error message, and
p8_clear_spans's trailing-padding span, which this geometry does not have; its test files predate the tail test inthis branch.
ruff 0.12.8(the version pinned in.pre-commit-config.yaml) checks and formats every Python file this PR touches.Precision notes. Cosines are computed in fp64. The RP2 direct rows (m <= 16) are graded against
the unquantized-activation reference; the grouped rows against the quantized one. The own-arms
session bounds every row at 0.9999, stricter than the grouped arm's 0.9998. Its m = 4096 row sits on
that bound: across the runs of this branch's trees it has printed 0.9998997 to 0.9999003, moving with
the race's winner and run-to-run rounding, so it passes or prints LOW from run to run; no bound was relaxed.
The kernel takes each UE8M0 exponent from the exact FP32 bits of
amax / 448and clamps it to[-127, 127], and the three CPU references implement the same rule: torch's float32
log2rounds somequotients just above a power of two down to that power, which would pick the smaller scale.
Related pull requests