Repository navigation
Alpamayo 1.5 on Thor: full-request engine with precision presets and FlashDrive streaming backend - #229
Open
LiangSu8899 wants to merge 21 commits into
Open
Alpamayo 1.5 on Thor: full-request engine with precision presets and FlashDrive streaming backend#229LiangSu8899 wants to merge 21 commits into
LiangSu8899 wants to merge 21 commits into
Conversation
The FlashDrive backend no longer calls vLLM Marlin, cuBLAS or torch elementwise ops inside the streaming windows: - W4A8 GEMMs for the verification block and the trajectory-start row run on a few-row mma.sync kernel over the packed INT4 nibbles, zero points and integer group scales (Marlin's integer arithmetic exactly; one warp per 16-row tile streaming the weights at the DRAM rate, K split across warps for narrow outputs; a shared-memory ring variant for N < 2048). - The DFlash draft (fc + hidden norm, both layers with FA4 attention of the block rows over context + block keys, SiLU MLP) and the LM head run on a few-row BF16 kernel over fragment-ordered weight copies; `decode_weights="fp8"` keeps per-row FP8 copies for the W8A16 kernel. - The vision tower and the action expert run on tcgen05 BF16 GEMMs with the bias, GELU and residual add fused into the epilogue, a LayerNorm kernel and a fused SiLU-mul over the gate|up output; `expert_weights= "fp8"` adds a W8A8 path (per-row FP8 activations and weights, scales in the epilogue). `vision_gemm`/`expert_gemm="torch"` keep cuBLAS. - The frame-block shift of the whole cache (keys, values, rotated view) is one launch. - Options: `act_rot="dense"` (the ParoQuant rotation as one 128x128 FP16 matrix per group on tensor cores), `w4a8_ktile=256`, `pdl` (programmatic dependent launch of every backend kernel, off by default). The W4A8 activation buffers are padded to whole 32-row tiles and the persistent buffers zero-initialised: the tcgen05 prefill GEMM reads the padded last token tile, which made two processes with the same seed diverge. Two processes now reproduce all 117 windows of the default clip. Tests cover exactness against float64 references and Marlin, the fragment layouts, repeatability of every new kernel and of every primitive on identical state, and the three backend configurations.
…ow in one kernel per input Over the streaming window the producer of each W4A8 input (residual update + RMSNorm, SiLU-mul of the gate|up output, or the attention output) and its ParoQuant preparation (channel scale, eight rotation layers, row absmax, INT8 quantization) now run as one kernel over tiles of 8 or 16 rows. The tile is held column-major in shared memory as 16-byte chunks of 8 rows (one pad chunk per 8 channels), so a rotation lane moves 8 rows per access and the producers' transposing stores stay conflict-free; the widest input (down projection) spreads a tile over a thread-block cluster whose CTAs share the row absmax through distributed shared memory. The RMSNorm keeps the summation order of the separate kernel, so every output is bit-identical to the previous path (unit tests over all producer modes, tile sizes and cluster shapes). The few-row verification and trajectory-start passes keep the separate kernels, whose grids spread the rotation groups over the SMs. Thor, author clip, seed 0, 117 windows: prefill 177.8 -> 155.6 ms, window 482.0 -> 465.6 ms mean (paired -16.4), tokens identical.
… K tile, INT8 weight tiers The pairs of every ParoQuant rotation layer and channel group are reordered on the host so that the eight pairs a quarter-warp of the tile preparation kernel handles together touch distinct shared-memory bank groups with their first channel; the rotation itself is unchanged (pairs of a layer are disjoint) and the outputs stay bit-identical. The blockwise tcgen05 prefill GEMM now runs two 128-channel groups per mainloop step by default (identical integers). Two precision tiers for the streaming path, both off by default: - decode_weights="int8": per-row symmetric INT8 copies of the LM head and the DFlash draft linears on the W8A16 kernel, which gains an INT8 code path next to FP8 E4M3 (same bytes, ~4x finer weight steps). Nine clips x three seeds on Thor: ADE +0.008 vs the BF16 default (95% CI [-0.031, +0.063]), -0.023 vs the FlashDrive author stack; block step 37.4 -> 30.6 ms per round. - prefill_weights="int8": per-channel INT8 requantization of Marlin's group-scaled integers on a plain INT8 tcgen05 GEMM with the token and channel scales in the epilogue (171-179 TOPS vs 115-138 blockwise on Thor). Thor, author clip, seed 0, 117 windows, defaults: window 465.6 -> 454.4 ms mean (paired -11.2), prefill 155.6 -> 144.4, tokens identical.
…sured tiers in the docs The per-channel INT8 prefill GEMM runs on 256 x 128 x 128 tiles with the 2-SM UMMA (2 x 1 clusters): 201-211 TOPS on Thor against 171-179 for the 1-SM 128 x 128 tile, identical integer results. The blockwise W4A8 GEMM gained the same 2-SM form behind its flags for comparison (no gain there: its per-group promotion, not operand traffic, is the limit) and a test that the K tile 128 / 256 and 1-SM / 2-SM forms produce identical outputs. Docs: nine clips x three seeds on Thor for the defaults and every tier (INT8 decode weights, INT8 prefill weights, FP8 expert weights, combinations), latency and trajectory error paired against the FlashDrive author stack.
…and prefill tiers Full-request engine (Alpamayo15Thor): - vision tower on the tcgen05 BF16 GEMMs with bias, GELU and residual fused into the epilogues (2-SM 256 x 128 tiles for every form), LayerNorm kernel, one CUDA graph per frame grid; vision_gemm="torch" keeps the cuBLAS path - action expert on 64-row tcgen05 tiles with fused q/k/v and gate|up weights and the split SiLU-mul kernel (no cuBLAS split-K reduce launches) - prefill: gate and up as two cuBLAS GEMMs over the fused weight (cuBLAS is slower on the fused shape), vectorized SiLU-mul in place; the tcgen05 BF16 GEMM measured slower than cuBLAS at these shapes and is not used there - precision tiers: decode_weights="int8" (W8A16 INT8 path), expert_weights="fp8" (W8A8 per-row), prefill_weights="fp8" (W8A8 per-row on 2-SM FP8 tiles; the K=12288 down projection as four column blocks that fit the L2) Kernels: FP8 GEMM gains 128-row and 2-SM 256-row tiles and an output leading dimension; the LayerNorm, split SiLU-mul (now gate/up pointers with a leading dimension) and per-row FP8 quantization kernels move under the alpamayo15_ prefix as shared Alpamayo 1.5 kernels; the quantization uses the same division as the torch reference. The FlashDrive backend imports the shared GEMM wrapper, graph capture and vision tower from the engine. Tests: shared kernel tests moved to the engine file, 2-SM tile identity, FP8 column-block outputs, option validation, a synthetic Qwen3-VL tower comparing the torch, CUTLASS and captured-graph paths.
…ducers The W8A8 prefill tier quantized the RMSNorm and SiLU-mul outputs in a second pass over each row. The producers now write the FP8 codes and row scales themselves (rms_fp8_rows, silu_mul_fp8_rows: same codes and scales as the BF16 kernel followed by the row quantization, the BF16 row is never written), leaving one quantization launch per layer for the attention output. Tests compare both fused kernels with the separate ones bit for bit.
…d the 9-clip measurements
… producers, PDL and kernel tuning - decode_weights="w4a8" with paro_checkpoint: the decode projections run as the ParoQuant W4A8 linears of the Alpamayo-1.5-10B-PARO checkpoint on the few-row W4A8 kernels (activation preparation + Marlin integer arithmetic), the LM head as INT8. flash_rt/models/alpamayo15/marlin_w4a8.py holds the Marlin layout index maps (following vLLM's marlin_utils, Apache-2.0), the unpacking, the few-row operand packing and the checkpoint loader; the FlashDrive backend imports them. The few-row W4A8 GEMMs gain an output leading dimension and a BF16 output. - vision_weights="fp8": W8A8 vision blocks on a 2-SM FP8 GEMM with bias, GELU and residual fused into the epilogue; LayerNorm -> FP8 rows kernel. - the W8A8 expert tier quantizes inside its FP32-residual RMSNorm and SiLU-mul producers; shared row-quantization device helper header. - pdl: programmatic dependent launch on every engine kernel and GEMM (alpamayo15_set_pdl, shared with the FlashDrive backend; off by default). - RMSNorm kernels issue four loads per step in the same per-thread order (bit-identical outputs); the vision RoPE kernel uses one thread per vector pair (no idle lanes). Tests: Marlin layout round trip, fused producers against the separate kernels bit for bit, FP8 GEMM epilogue forms, LayerNorm -> FP8 rows, option checks.
…nts of the W4A8 configurations
… and table prefetch, INT4 LM head Engine presets bf16 (default), balanced (W4A8 decode + FP8 expert + FP8 prefill) and fast (balanced + INT4 LM head); the weight options override a preset. Decode: the ParoQuant projections of a shared input run as one grouped few-row W4A8 GEMM launch (q|k|v, o, gate|up, down) on a weight layout that carries each tile's zero points and integer group scales in front of its nibbles; up to four rows the per-token INT8 quantization runs in the GEMM prologue (rotation pass + GEMM per group), bit-exact with the two-pass preparation; each GEMM prefetches the next group's rotation tables into the L2 while streaming its weights (a cold read of a few hundred KB right before a streaming kernel costs ~30% of that kernel on Thor). FA4 decode attention with one KV split. Own INT4 LM head (group 128, zero points, per-channel normalized integer scales rounded up, padded vocabulary, per-token INT8 activations). Vision LayerNorm one warp per row (also the LayerNorm -> FP8 producer); the prefill RMSNorm -> FP8 producer reads and writes 8-byte vectors. Author clip, 24 windows: balanced 895 / 788 ms (decode 22.1 ms/token), fast 849 / 773 (20.8); nine clips: balanced 1027 / 972, fast 1006 / 949 (ADE neutral or better against the default and the official model).
…t default's token identity
…measured backends and tiers A cuBLASLt BF16 GEMM with bias, GELU and beta-1 residual fused into its epilogue (column-major formulation, cached plans) and vision_gemm="cublaslt", which runs the blocks' fc1 on it: 147 against 125 TFLOPS on Thor, about 3 ms per request, but last-bit differences that resample a few windows (nine clips: ADE +0.047 [-0.007, +0.153]), so the tcgen05 tower stays the default. The other block GEMMs measured no better on cuBLASLt in the pipeline (the fc2 choice 1.35 ms against 0.90). The prefill tier's SiLU-mul -> FP8 producer keeps four vectors in flight (same codes). Measured and left out: the prompt's last-token logits from the quantized LM head (5 ms, ADE +0.055 [+0.001, +0.161]) and the repository's NVFP4 GEMM at the prefill shapes (224 TOPS against 250 for the FP8 2-SM tile).
…t preset takes it NVFP4 GEMM forms for the prefill shapes in the flash_rt_fp4 module (1-SM 128x128x256, 2-SM 256x128x256 and 256x256x256 block-scaled tiles; on Thor the 256x256 2-SM tile runs 500-575 TOPS against the FP8 tile's 250), with the output leading dimension and a device alpha. Row quantizers for bf16 operands: a per-tensor scale from per-block maxima that keeps the UE4M3 block-16 scales in the normal range, e2m1 codes by the hardware conversion, scale factors in the CUTLASS tile-interleaved layout, the GEMM alphas of an input's consumers. prefill_weights="nvfp4": static NVFP4 weights, each GEMM input quantized once per consumer group, the 36 layers' projections on the 2-SM tile. Prefill 290 -> 233 ms per request on the author clip; nine clips 884 / 836 ms (3.76x the official model) at a neutral ADE (-0.022 vs the BF16 default [-0.090, +0.026], -0.008 vs official [-0.076, +0.054]). The fast preset is now balanced + NVFP4 prefill + INT4 LM head.
…-row alpha GEMM epilogue The W4A4 prefill tier quantized each GEMM input with a per-tensor scale in two passes (per-block maxima, then codes) after a BF16 producer. The activations now carry per-row scales, so the RMSNorm and SiLU-mul producers quantize their row completely from shared memory (codes, UE4M3 block scales in the GEMM's scale layout, and the row's alphas row scale x weight scale for the input's consumers), the attention output goes through a per-row quantizer, and the NVFP4 GEMM applies the alpha vector in its epilogue (PerRowLinCombPerRowBias with beta and bias off; the scalar form stays for the weights). The BF16 intermediates and both passes disappear from the tier: prefill 233 -> 185 ms, the fast preset 807 -> 741 ms on the author clip and 884 -> 837 ms on the nine clips (ADE -0.036 vs the BF16 default [-0.139, +0.036]). Remove a copy of the prefill's NVFP4 block that had been left in the expert forward. Tests: per-row quantizer and per-row alpha GEMM against the torch emulation (zero rows, second consumer, exact per-row powers of two, small rows no worse than large), fused producers against the separate kernels plus the per-row quantizer (residual bit-exact, SiLU bit-exact, RMSNorm within one e2m1 step), the gate | up stride form, wider rows, rejections.
…EMMs Add tile 3 = 256 x 256 x 64 on a 2-SM UMMA (identity, GELU tanh / erf and residual epilogues) to the BF16 GEMM and make it the vision tower's default: at S = 11520 qkv 736 -> 551 us, fc1 972 -> 708, fc2 969 -> 816, proj 331 -> 337 (N = 1152). The N tiling leaves the K reduction order unchanged, so the outputs equal the 256 x 128 tile's on every GEMM and the tokens and trajectories are the same; the tower drops from 141 to 122 ms per request.
…dow default and fast rows
preset="balanced" (W4A8 ParoQuant decode, FP8 expert, FP8 prefill) becomes the default: over three seeds of the nine-clip evaluation it is neutral against the official model (ADE -0.001 [-0.052, +0.053]) at 3.4x its speed; on the author clip 867 / 870 ms over 120 windows against 1873 for the BF16 engine. It needs the ParoQuant checkpoint directory (paro_checkpoint); the engine refuses to start without it and names preset="bf16" as the official-weights alternative. The option checks run before the checkpoint check, so a wrong option is reported as such.
…takes presets and tiers Rewrite docs/alpamayo15.md around the two inference forms (full requests against the official model, streaming against FlashDrive): installation, checkpoints, usage, one benchmark protocol (nine clips, three seeds, model time per window, paired ADE with clip-level intervals) and one results table per form, each arm against its own reference, plus a stage table per form on the example clip; precision tiers and execution contracts follow. The historical per-step tables and the earlier eight-clip protocol are dropped. benchmarks/alpamayo15_thor.py gains --preset (bf16 / balanced / fast; the engine's default is balanced, so the runner resolves --paro into the ParoQuant directory) and --engine-kw for single tiers of the engine or the streaming backend.
LiangSu8899
changed the title
…esets, streaming backend), FlashDrive acknowledgment
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.
Related issue: #223
Alpamayo 1.5 on Jetson AGX Thor (SM110) in two opt-in forms, each measured against its own reference on the same inputs:
flashrt_full): the official request procedure (nvidia/Alpamayo-1.5-10B) on FlashRT kernels, with three precision presets. Reference: the official model.fd_flashrt): FlashDrive's streaming procedure (z-lab base / PARO / DFlash checkpoints, FlashDrive's orchestration) on FlashRT backends. Reference: FlashDrive's own implementation.The two procedures produce different outputs and are not compared with each other.
Performance
Thor, locked clocks, NVIDIA PyTorch 26.05 / CUDA 13.2. Nine PhysicalAI-AV clips, three seeds, 1719 steady windows per arm, 128 reasoning tokens, one trajectory sample; latency is model time per window (tokenized inputs on the host to the trajectory on the host). Paired ADE differences carry 95% bootstrap intervals over clips.
Full requests:
preset="bf16"preset="balanced"(default)preset="fast"Streaming:
torch.compile)Every ADE interval includes zero: each FlashRT path is indistinguishable from its reference on this evaluation (offline ADE of sampled trajectories; not a closed-loop metric).
Presets and tiers
bf16balanced(default)paro_checkpoint(z-lab PARO)fastparo_checkpoint,flash_rt_fp4moduleStreaming backend tiers (off by default, each validated above and singly in the guide):
decode_weights="int8",prefill_weights="int8",expert_weights="fp8".Usage
Python APIs:
Alpamayo15TorchFrontendThor(checkpoint, engine_kwargs={"paro_checkpoint": ...})andAlpamayo15FlashDriveTorchFrontendThor(FlashDriveCheckpoints(...)). Installation, pinned dependencies, checkpoints, the protocol, stage breakdowns, numerics contracts and limitations: docs/alpamayo15.md.Validation
On Thor at the head commit:
tests/test_alpamayo15_thor.py131 passed,tests/test_alpamayo15_flashdrive_thor.py66 passed (kernel exactness against the official op sequences and against Marlin / ParoQuant, fused-vs-separate kernel identity, preset and option validation, checkpoint integration). Thebf16engine keeps the official numerics up to kernel reduction order;balanced/fastand the streaming tiers are numerics changes gated by the nine-clip, three-seed evaluation above. Kernels are Thor-only and default OFF (FLASHRT_ENABLE_ALPAMAYO15); model dependencies stay optional.