diff --git a/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md b/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md new file mode 100644 index 0000000000..975939075a --- /dev/null +++ b/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md @@ -0,0 +1,567 @@ +# Flash-Next prefill: 6K investigation + +Integration baseline: `onecat/main`, `6f3557c973b4602905037e086fd4acba9b0d867f` +(PR #1203). Work is isolated on +`codex/v100-flashnext-prefill-6k-9683-20261013-160548`. + +## Workload and evidence boundaries + +- Four V100-SXM2 16 GB GPUs, NV2 links, CUDA 12.8, driver 570.211.01, + Torch 2.10.0+cu128, Python 3.12.3. One 18-core/36-thread E5-2686 v4 socket. +- Flash-Next IQ3_S, TP4, MTP4, maximum context **262144 tokens**, + maximum scheduled chunk **8192 tokens**, prefix caching and thinking enabled. + Experts retain 80% GPU / 20% host residency; QSA KV stays on host, PLE on disk. +- Production sampling remains temperature 0.6 / top-p 0.95 / top-k 20. + The historical `ttft.py` is an explicitly greedy, one-output-token timing test. +- Historical `ttft.py 32000` uses random words, approximately 34700 tokens, + **five** chunks (7344 x 4 plus approximately 5347). Its 7.50–7.55 s figure + measures a complete HTTP response, not streaming TTFT or pure prefill. +- Historical `profile_request.py` recreates its RNG inside the word generator, + producing one repeated word and **four** chunks. Its 6.48 s trace cannot close + the wall-clock accounting for the different 7.5 s request. The old table also + omitted 0.91 s of other kernels when only its first five rows were summed. +- For approximately 34700 tokens, 6000 tokens/s corresponds to approximately + **5.8 s**: approximately 1.7 s / 23% below the historical response latency. + No new end-to-end speed claim is established by the microbenchmarks below. + +## Resource findings from the existing DMA trace + +The critical GPU's main-stream kernel/copy union occupies 87.35% of the +6.482 s trace interval. Across all streams, kernel union occupies 88.98% and +kernel/copy union 97.56%. Main-stream gaps are not whole-device idle time. +Cold-expert copies overlap kernels for 85–87% of their duration; their summed +3.45–3.70 s per GPU must not be added to end-to-end latency. + +Matching collective ordinals across all four ranks gives 400 all-reduces +and 812 all-gathers per rank. GPU0 arrives last at 278/400 all-reduces. +The sum from the latest rank arrival to the latest completion is 400.8 ms +for all-reduce, compared with GPU0's 419.5 ms kernel sum. The earliest-to-latest +arrival spreads sum to 258.6 ms; those waits mostly occur on earlier ranks. +They cannot simply be subtracted from the critical rank's latency. These are +trace timing estimates, not link bandwidth counters or a new workload result. + +The original analysis joined CUDA correlation IDs across TP processes. +`nsys_breakdown.py` now keys them by process and correlation ID, and keys the +main stream by process and stream ID. Replaying the existing DMA trace moves +20 ms from the host-gap bucket to GPU waits: **133 ms host gaps, 687 ms GPU +waits**, with the same 6482 ms span and kernel totals. These are gap +attributions, not hardware utilization counters. + +The trace has launch geometry and registers/shared-memory metadata, but no SM, +Tensor Core, DRAM, or NVLink hardware utilization counters. Static occupancy +limits are not measured resource saturation. For example, QSA rows use +120 registers/thread and approximately 34 KiB shared memory per block, allowing +two 256-thread blocks/SM. Increasing this to three was already rejected in the +previous investigation. + +The saved real QSA input has 8160 queries, six heads, dimension 256, and 2051 +index positions per query. Only 12240 of 16736160 indices are negative. +Reading each valid K/V pair once would request approximately 8.56 GB; dividing +by kernel time is **requested payload bandwidth**, not measured HBM traffic +(cache hits are not known). + +## Experiments and decision gates + +### 1. Narrow grouped MoE tiles + +Hypothesis: expert batches contain roughly 100–160 rows at the current chunks. +Additional 64x64 tiles may reduce masked work, accumulator registers, and N=320 +padding. Reuse the existing TurboMind grouped provider, codec transforms, +autotuner, workspace, and group-sized K iteration. Existing kernels remain +available as fallback. No graph, cache-state, or sampling behavior is changed. + +The experimental tiles use the existing narrow-prefill admission +(`M >= 512` routed rows, `N <= 512`, `N % 64 == 0`). Existing provider validation +still applies. Measured cases are much narrower: TP4 rank-zero real weights, +512 experts, top-10, hidden 2560, intermediate 160, N=320 gate/up; down executes +the production 4096-token sub-batches and unroute operation. This does not +qualify every admitted shape or other models. + +`moe_tiles.py` tests all gate/up types present in this model: +18 (17 layers), 22 (20 layers), 21 (10 layers), 23 (one layer). Down uses +20 (39 layers) or 42 (nine layers). Timing excludes routing, transcode, +host DMA, and multi-GPU communication. It includes both GEMMs, SiLU, down +sub-batching, allocation calls, and output unroute. + +Predeclared microbenchmark numerical gates: finite hidden/output tensors; +maximum checked-row relative L2 <= 0.01 against independently dequantized +FP32 gate/up and complete MoE references; candidate/control checked-row +relative L2 <= 0.003. These tolerate FP16 accumulation/rounding, but do not +replace paired model task-quality, repetition/EOS, or long-context checks. +Require a stable chain improvement before spending a service restart on it. + +Baseline, median of five CUDA-event samples, five iterations/sample: + +| Layer / gate type | Tokens | Gate/up + SiLU ms | Down + unroute ms | Complete chain ms | +|---|---:|---:|---:|---:| +| 0 / 18 | 7344 | 4.327 | 3.156 | 7.367 | +| 0 / 18 | 5347 | 3.759 | 2.294 | 5.628 | +| 1 / 22 | 7344 | 4.042 | 2.908 | 6.952 | +| 1 / 22 | 5347 | 2.482 | 2.061 | 4.506 | +| 17 / 21 | 7344 | 4.458 | 3.150 | 7.465 | +| 17 / 21 | 5347 | 3.373 | 2.280 | 5.700 | +| 47 / 23 | 7344 | 3.875 | 3.200 | 6.833 | +| 47 / 23 | 5347 | 2.924 | 2.323 | 5.175 | + +All eight baseline cases pass; maximum checked output error is 0.000844. +Native dispatch confirms type-18 currently selects 32x128, so a benefit cannot +be inferred merely from reducing the M tile. Separate-stage times should not +be summed as a replacement for the measured chain (GPU clocks vary). + +Rejected before GPU testing: 64x64 with four warps cannot instantiate the +current metadata thread map (zero-width access for this layout). Retain the +two-warp candidate for measurement; do not change shared iterator contracts +just to admit the extra tuning candidate. + +A second artifact adds padded activation supply using the existing quantized +batch-family operand. Hypothesis: at K32/K64, the unpadded shared-memory row +stride repeats banks across the distinct rows of each eight-lane MMA load +group. Padding by eight half values separates those accesses. This is a +static layout hypothesis, not hardware-counter evidence. Candidates are +32x128x32 with 32x64 epilogue subtiles and 64x64x64 with four warps. The K64 +metadata tile has enough elements for four warps. Decoder, group size, +routing, bank storage, and fallback kernels remain shared with the existing +family. Separate artifacts retain the first candidate for comparison. + +Measured complete MoE chains (ms): + +| Layer / tokens | Control | 64x64 | Padded candidates | Padded reduction | +|---|---:|---:|---:|---:| +| 0 / 7344 | 7.367 | 7.047 | 6.995 | 5.1% | +| 0 / 5347 | 5.628 | 5.124 | 4.956 | 11.9% | +| 1 / 7344 | 6.952 | 6.868 | 6.318 | 9.1% | +| 1 / 5347 | 4.506 | 4.537 | 4.404 | 2.3% | +| 17 / 7344 | 7.465 | 7.224 | 6.711 | 10.1% | +| 17 / 5347 | 5.700 | 5.172 | 5.006 | 12.2% | +| 47 / 7344 | 6.833 | 6.736 | 6.329 | 7.4% | +| 47 / 5347 | 5.175 | 4.808 | 4.507 | 12.9% | + +All eight cases pass for both candidates. Checked hidden/output rows are +bit-exact to control. Weighting these isolated chain deltas by 48 layer types +and five chunks gives only about 0.13 s for the padded artifact; this is a +screening estimate, not an end-to-end result. No claim of 6K follows. +The added grouped-prefill regression covers M=511/512/513, interior M=2049, +N=64/320/512, N32 and wide fallbacks, K=256/512/768/2560, uneven/empty experts, +indexed and direct input, fused SiLU, output guards, and changed-input graph +replay. Its predeclared maximum-row relative-L2 limit is 0.005 versus an +independent GGUF-dequantized FP32 reference. + +**The second padded artifact is rejected despite those eight passes.** The +boundary regression finds non-finite/wrong results in five of 24 cases. +Its 32x128 CTA used a 32x64 epilogue subtile to distinguish the tuning key, +but that subtile is narrower than the output thread map's horizontal span. +Some threads therefore read outside the intended shared-memory tile. The +third candidate removes it and uses a full 32x64x64 CTA/epilogue with a 2x2 +warp map, alongside the full 64x64x64 candidate. No shared epilogue or +iterator contract is changed. A 1x4 warp map was rejected at compilation: +its N=16/warp is below this MMA family's N=32 atom. + +The verified original artifact passes **24/24**. Pytest initially prepended +the candidate source root through the tests package, so a nominal control +invocation actually loaded the padded artifact; that comparison is invalid +and retained only as a harness failure. The corrected runner imports vLLM +and `_C` before invoking pytest, asserts their paths against the requested +artifact, and prints the native SHA256. It confirms 24 original passes and +five failures for the second padded artifact. Microbenchmark JSON now also +records loaded package/native paths and native hash. + +The corrected third artifact passes **24/24 boundary cases**. Its full-chain +times for the eight rows above are 6.723, 4.753, 6.201, 4.437, 6.672, 4.824, +6.404, and 4.592 ms: reductions of 1.5–15.5%, with all checked hidden/output +rows bit-exact to control. The layer/chunk-weighted isolated delta is 0.163 s. +This artifact, rather than the rejected second one, proceeds to model testing. +Its `_C` SHA256 is +`5d389732b5fb828b947ffdb74a3ec6c80aa70cb60bdc6b1ca65e437a9514148e`. + +Two smaller hypotheses were rejected on the original artifact: + +- Increasing down sub-batches from 4096 to 8192 tokens: layers 0/1 at + 7344/5347 tokens give complete-chain times of 7.590/5.921 and 7.269/4.878 ms, + versus 7.367/5.628 and 6.952/4.506 ms. All checked outputs are bit-exact to + control; speed is 3–8% worse, with a larger temporary output allocation. +- Tuning gate/up at the actual 7344 rows rather than 8192: complete chains + are 7.378 and 6.908 ms for layers 0/1, versus 7.367 and 6.952 ms. There is + no stable material improvement; checked outputs are bit-exact to control. + +### 2. HC narrow projection + +Hypothesis: the M~8K, N=88, K=10240 HC projection exposes too few concurrent +blocks. `hc_projection.py` compares the existing cuBLAS call with two, four, +and eight K partitions, FP32 partials and FP32 reduction. It includes strided +batch handling and reduction costs. Synthetic data screens throughput only; +any admitted candidate needs real activations/weights and paired quality. +The numerical screen is finite output and maximum row relative L2 <= 0.001 +against FP32 and against the current projection. + +Rejected: for 7344 rows, control is 0.503 ms and K partitions 2/4/8 are +0.532/0.537/0.568 ms; for 5347 rows, control is 0.360 ms and partitions are +0.425/0.419/0.411 ms. All outputs are finite; maximum relative L2 is 0.000298 +against FP32 and 0.000602 against control. More parallel blocks do not offset +the partial-output/reduction overhead. + +### 3. Larger custom all-reduce admission + +`collectives.py` uses all four idle GPUs, rank-max CUDA-event times, FP16 +rows x 2560 messages, and the production PyNccl/CustomAllreduce providers. +It includes eager custom-AR staging into its IPC buffer. Raising its capacity +to 48 MiB adds 80 MiB of GPU buffers per rank relative to the default 8 MiB. + +| Rows | NCCL AR ms | Custom AR ms | NCCL AG, width 88 ms | NCCL AG, width 640 ms | +|---|---:|---:|---:|---:| +| 7344 | 0.547 | 0.585 | 0.084 | 0.318 | +| 5347 | 0.411 | 0.431 | 0.076 | 0.254 | + +Rejected: increasing custom-AR capacity alone is 5–7% slower. Checked rank-zero +outputs are finite, relative L2 <= 0.000370. This is a synthetic throughput +screen, not full collective correctness or model qualification. Isolated NCCL +is faster than historical in-service call averages: rank arrival skew and +concurrent resource use need to be separated from transfer cost before another +communication change. + +### 4. Existing QSA rows route + +The saved real prefill input confirms the current rows route at 12.54 ms, +versus 23.07 ms for the old tiled route (1.84x). Maximum absolute output error +is 0.000977 and maximum per-head relative L2 is 0.000774; no NaNs. Requested +payload divided by runtime is about 683 GB/s, not a measured HBM counter. +This gain is already in the baseline and must not be counted again. + +### 5. HC communication/compute pipeline + +`hc_pipeline.py` uses the existing HC packet, gate-mix, and assembly kernels +with cuBLAS and PyNccl. It splits rows into two/four chunks and overlaps the +two gathers with subsequent projection work on a side stream. The timing +includes submissions, events, assembly, and final concatenation. + +Rejected: 7344-row sequential time is 1.343 ms, versus 2.390/4.985 ms for +two/four chunks. At 5347 rows, times are 1.056 versus 2.642/4.913 ms. This +Python pipeline adds more overhead than it hides. Output versus the same +tiles executed serially is bit-exact, isolating stream ordering from cuBLAS +shape-dependent arithmetic. The original maximum-row relative-L2 screen +of 0.001 also fails for the four-element injection output in three cases +(maximum 0.004036); it is not relaxed. Model-output errors are <= 0.0000761 +and all outputs are finite. No runtime change or quality qualification follows. + +### 6. Per-projection prefill streams + +Hypothesis: the shared cold stream can hold down DMA behind a gate/up bank +wait, and the shared resident stream orders next-layer gate/up transcode +before current-layer down transcode. `moe_streams.py` screens separate cold +streams and separate cold/resident streams independently, using the actual +native expert cache with 402 resident experts plus eight spare slots and +136 host slots. Canonical banks and staging bytes do not increase. + +Eight logical layers repeat four real TP4-rank-zero layer formats, wrapping +the four-view ring. Output and restarted-pass checks must be bit-exact; +native decode misses then change residency before another comparison. The +split-resident experiment explicitly joins each view's down readers before +reusing it, including leftover prefetch after an interrupted pass. Optional +GPU sleep models only a between-layer delay, not model compute contention. +This remains a benchmark-only hypothesis pending measured benefit. + +Rejected: with no artificial gap, eight-layer pass medians are 130.18 ms +shared, 130.59 ms cold-only split, 127.23 ms fully split, and 130.86 ms on +the repeated shared control. With a 20-million-cycle layer gap, all modes +take 186.3–186.9 ms. The small isolated benefit disappears once between-layer +work hides transfers. All output/restart comparisons are bit-exact, including +after 256 expert loads, 208 evictions/write-backs, and zero cache errors. +There is no runtime stream change. + +### 7. Reuse the existing dequantize-plus-cuBLAS dense provider + +The old trace attributes 450 ms to Q6_K bit-plane dense kernels and 149 ms to +Q5_K, across four chunks. They use 128x256 CTAs, approximately 248–255 +registers/thread, and one resident CTA/SM. These are static resource limits; +no Tensor Core utilization counter was collected. Dispatch inspection confirms +the affine fused route uses the default policy for M > 32, so an 8192-row +warmup does not autotune the measured 7344/5347 shapes. + +`dense_blas.py` uses real rank-zero TP4 weights and includes weight +dequantization into temporary scratch on every call. It compares fused, +dequantize+BLAS, then fused again. Maximum row-relative L2 against the FP32 +reference and paired output must be <= 0.003, with finite outputs. All ten +large-M cases and 25 threshold-screen cases pass. + +| Projection / source | M | Fused ms | Dequantize+BLAS ms | Fused repeat ms | +|---|---:|---:|---:|---:| +| QKV / Q6_K | 7344 | 1.867 | 1.191 | 1.839 | +| QKV / Q6_K | 5347 | 1.370 | 0.917 | 1.362 | +| QKV / Q5_K | 7344 | 1.908 | 1.143 | 1.921 | +| QKV / Q5_K | 5347 | 1.411 | 0.900 | 1.408 | +| SSM output / Q6_K | 7344 | 1.024 | 0.711 | 1.049 | +| SSM output / Q6_K | 5347 | 0.797 | 0.563 | 0.795 | +| Attention gate / Q6_K | 7344 | 1.929 | 0.697 | 1.975 | +| Attention gate / Q6_K | 5347 | 0.979 | 0.587 | 0.975 | + +Only four physical layouts are added to the existing calibration table: +(bits, group, K, N) = (5,32,2560,2560), (6,16,2560,2560), +(6,16,1536,2560), and (6,16,2560,1536), all at **M >= 2048 tokens**. +At M=512/1024 several layouts regress, so those retain fused computation. +The N=160 shared-expert screen is inconsistent and is not admitted. +Uncalibrated layouts, missing native operators, and failed workspace allocation +retain the existing fused fallback. No codec, provider, startup flag, or +workspace owner is added. The existing shared maximum-size scratch is retained; +its model-level memory impact is part of the pending service check. + +The affected GPU suite passes **15 tests**, including actual native route +observations at M=2047/2048/2049/4097, both eager and compiled calls, the +serving opaque projection, shared scratch, and changed-input graph replay. +Correctness support remains the existing affine provider's valid FP16 matrix +contract. Acceleration qualification is those four layouts; measured speed +points are M=256/512/1024/2048/4096/5347/7344, not every possible M or TP rank. + +### 8. Skip cold-expert holes while preserving staging offsets + +After decode, 110 cold experts occupy an extent of 135–136 host slots. The +old staging copy transfers that whole extent, including slots that the cold +transcode never reads. `cold_dma.py` first screens fragmented 110/104-slot +maps at real shard sizes: full copies of 73.98 MB take 16.05–16.37 ms; +selected copies take 13.31/12.71 ms including Python/native submissions. + +`ExpertCacheLane.stage_cold` now derives contiguous occupied ranges from the +same saved residency map consumed by the transcode, keeping every destination +offset unchanged. It reuses the existing host arena, staging buffer, stream +dependencies, and cache owner. It does not change residency, arithmetic, +spare slots, or GPU memory. A host slot must be valid for every cold expert; +invalid mappings raise an explicit error. + +`moe_dma.py` compares both copy policies in the actual eight-layer prefill +pipeline, with all other operations identical. Times are medians in ms: + +| Cache state / between-layer gap | Full extent | Occupied ranges | Full repeat | +|---|---:|---:|---:| +| Initial contiguous / none | 132.857 | 133.717 | 134.331 | +| Initial contiguous / 20M GPU cycles | 181.523 | 182.574 | 182.467 | +| Fragmented after decode / none | 161.027 | 136.131 | 162.643 | +| Fragmented after decode / 20M GPU cycles | 187.707 | 182.438 | 187.692 | + +All outputs and interrupted-pass restarts are bit-exact, including after 888 +expert loads, evictions, and write-backs, with zero cache errors. The gap is +synthetic delay, not real model compute. The complete affected cache suite +passes **9 tests**, including graph decode, a single hardware queue, mixed +prefill/decode, fragmented/slack-free host arenas, untouched staging holes, +and using the saved snapshot after the live CPU residency table changes. +Full-model latency remains a separate acceptance gate. + +### 9. Reject padded expert-batch FP16 BLAS + +`moe_blas.py` screens real layer-zero TP4 weights at 7344 tokens, top-10, +512 experts and 192 padded rows/expert. It includes per-batch dequantization, +padded activation gather, gate/up, SiLU, down, scatter and weighted unroute. +Contiguous compressed GPU banks are an optimistic input layout; production +would also need to gather resident/cold slots. Both arms omit routing and DMA. + +The current quantized chain takes 6.414 ms (repeat 6.478). Batches of 32, 64, +and 128 experts take 19.860, 17.301 and 16.972 ms respectively. All outputs +are finite with paired maximum row-relative L2 0.001142, within the preset +0.003 bound, but performance rejects this implementation. For 64 experts, +gate/up and down GEMMs alone take 0.688 and 0.255 ms, or about 7.54 ms for +all eight batches before dequantization/gather/scatter. FP16 gate/up output +only reduces its 0.688 ms to 0.661; it cannot close the gap. No runtime route +is added. Raw artifacts: `moe-blas.json` and `moe-blas-first.json`. + +### 10. Restore the calibrated policy after segment-bank restoration + +The matched serving trace exposed a missing integration path: no affine BLAS +or affine dequantization kernels execute, despite the four admitted physical +layouts. With `small_m_hmma=true`, initialization replaces canonical weights +with segment storage. Its prefill fallback restores canonical weights but +always invokes fused GEMM, and serialization discards the BLAS bands. +Thus the first 7.099 s result below does not demonstrate dense BLAS execution. +The earlier direct-kernel and ordinary opaque-path tests did not cover this +production storage mode. + +The follow-up carries the existing admitted bands through single and mixed +segment wrappers, retaining the small-M route. After canonical restoration, +it reuses the existing affine BLAS operator and engine-owned scratch only +inside those bands; all other inputs retain fused GEMM. Mixed output views +use a contiguous temporary because the existing native BLAS API requires it. +No native operator, configuration flag or extra persistent workspace is added. +`dense_segments.py` measures this actual restoration chain on real weights; +`test_gguf_segment_prefill.py` covers threshold neighbors, a range interior, +return to decode, compiled opaque calls, changed-input graph replay, mixed +output views and an uncalibrated fallback. **44 GPU/dispatch tests pass**, +including the existing segment codec, restoration, small-M graph and runtime +opacity regressions. Actual native BLAS dispatch is observed in the restored +chain; opaque and compiled surfaces produce the same qualified outputs. + +Real-weight restored-chain medians (ms), including canonical restoration and +weight dequantization on every call: + +| Projection / source | M | Restored fused | Restored BLAS | Fused repeat | +|---|---:|---:|---:|---:| +| blk.0.attn_qkv.weight / 14 | 7344 | 1.959 | 1.309 | 1.938 | +| blk.0.attn_qkv.weight / 14 | 5347 | 1.464 | 1.016 | 1.398 | +| blk.1.attn_qkv.weight / 13 | 7344 | 1.991 | 1.229 | 1.918 | +| blk.1.attn_qkv.weight / 13 | 5347 | 1.440 | 0.980 | 1.475 | +| blk.0.ssm_out.weight / 14 | 7344 | 1.148 | 0.824 | 1.144 | +| blk.0.ssm_out.weight / 14 | 5347 | 0.882 | 0.610 | 0.882 | +| blk.1.attn_gate.weight / 14 | 7344 | 2.077 | 0.814 | 2.038 | +| blk.1.attn_gate.weight / 14 | 5347 | 1.061 | 0.777 | 1.063 | + +All eight cases are finite and meet the preset 0.003 maximum row-relative +L2 bounds against both the FP32 oracle and the paired fused output. Full-model +route observation, quality and unprofiled speed with this fix remain pending. + +### Matched request measurement + +`measure_prefill.py` uses deterministic word sequences and per-request hashes, +clears prefix caching on an exclusively owned server, and records streaming +first choice, first text, complete response, actual token count, and service +metrics. Service metric sums are attributed only when their count increases +by exactly one. Server prefill includes scheduled wall time, not just GPU +kernel time. `--profile` uses identical request generation; profiled wall time +is not an unprofiled speed result. + +Fresh baseline on the owned server, after two excluded warmups (13.35/8.01 s): + +| Prompt tokens | Client first text s | Server prefill s | Server TTFT s | +|---:|---:|---:|---:| +| 34688 | 7.895 | 7.544 | 7.757 | +| 34606 | 7.513 | 7.186 | 7.448 | +| 34631 | 7.429 | 7.165 | 7.376 | + +Median client TTFT is **7.513 s / 4605.9 tokens/s**. In the middle request, +server TTFT exceeds server prefill by 0.261 s, and client first text follows +server TTFT by 0.066 s. These metrics share one request, unlike the historical +four-chunk profile. The first measured request still shows warming drift. +The existing 30K-token retrieval passes 2/2. `quality.py` adds paired archive +retrieval plus arithmetic, production T0.6/top-p0.95/top-k20 with thinking, +strict expected JSON, normal EOS, and a repetition check. The 32761-token +baseline passes with 166 generated tokens; near-256K qualification is pending. + +The first near-256K baseline was interrupted before completion to keep the +cheap optimization screen first. Above the configured 131072-token QSA +history-stage limit, the original runtime falls back to direct host-history +attention and slows sharply (around 32 layer snapshots per 30 s, all GPUs +busy but around 85–90 W). No long-context quality result or candidate +regression is inferred. Complete this baseline/candidate boundary check +before promotion; the configured limit and deployment remain unchanged. + +First candidate (before the segment-policy fix), same five prompt hashes and +token counts as control: + +| Prompt tokens | Client first text s | Server prefill s | Server TTFT s | +|---:|---:|---:|---:| +| 34688 | 7.502 | 7.171 | 7.387 | +| 34606 | 7.091 | 6.823 | 7.034 | +| 34631 | 7.099 | 6.790 | 7.021 | + +Median TTFT is **7.099 s / 4878.1 tokens/s**, improving 0.414 s (5.5%) from +the paired 7.513 s baseline. Excluded warmups take 14.13 and 7.56 s. +This is unprofiled; microbenchmark deltas are not substituted for this result. +It does **not** meet 6K tokens/s. The 32K quality prompt has the same hash and +passes the same expected JSON, EOS and repetition checks (167 generated tokens +versus control's 166). The 131070-token candidate also passes (161 generated +tokens, 60.44 s complete response); a matching 128K control is still pending, +so this does not establish a long-context performance comparison. + +The isolated cold-cache candidate startup took about 22 minutes, including +SM70+SM75 FlashQLA fallback compilation, checkpoint loading, and fresh main/MTP +graph compilation. This is not the warmed deployment's startup time. No OOM +occurred. The owned candidate was stopped after these checks for a bounded +expert-batch FP16 MoE screen and a matched Nsight profile. Its descendants +were recorded and only its verified API process received SIGTERM. + +### Matched candidate trace and boundary quality + +The first-candidate profile uses the same index-2 request as the unprofiled +runs: 34688 tokens, hash +`126bc5319be9170da4f07bd2bca263e083d2d636da17d0697d7a19d746b04342`. +Client first text is 7.312 s, server prefill 6.952 s and server TTFT 7.209 s. +The four main-stream spans are 6.914 s. This capture also enables NCCL +INIT/TUNING logging, so its wall time is not an unprofiled performance result. +The profiled prefill chunk row counts are 7344, 6528 and 6128; do not assume +that every request uses the historical 7344/5347 microbenchmark shapes. + +GPU0 closes as 6045.9 ms main-stream kernel/copy union, approximately 793 ms submitted +GPU wait and 75 ms host gap. The 397.9 ms outside the stream span remains +an unattributed HTTP/server residual; it is not added to kernel categories. +Main-stream categories from `nsys_breakdown.py`, in ms: + +| Work | Kernel/copy time | +|---|---:| +| TurboMind dense projections | 987 | +| MoE gate/up + down | 662 + 501 | +| MoE unroute + routing | 122 + 68 | +| QSA sparse partial + other + merge | 624 + 354 + 12 | +| AllReduce + AllGather | 501 + 381 | +| HC combine/norm + other HC | 348 + 125 | +| Other cuBLAS, including HC down | 589 | +| GDN | 350 | +| Elementwise/other, MTP, memcpy, dense helpers | 258 + 66 + 55 + 42 | + +The HC down cuBLAS kernel alone is 272.2 ms / 500 launches. GPU0 union +coverage is 87.44% main-stream kernel/copy, 89.88% all kernels, and 97.88% +all kernels/copies. These are timeline coverage, not SM, Tensor Core or HBM +utilization counters. NCU counters remain unavailable. Cold-stream H2D is +13.620 GB/rank over 3.553–3.766 s, about 3.62–3.83 GB/s. It overlaps +computation and must not be added to the wall-clock table. + +Across ranks, 500 main FP16 all-reduces spend 434.5 ms after the last rank's +arrival; arrival spread sums to 245.1 ms, with GPU0 arriving last 366 times. +The 1030 all-gathers have 335.9 ms after the last arrival and 133.3 ms spread. +NCCL 2.27.5 TUNING logs explicitly select RING/SIMPLE for the prefill payloads +(37,601,280-byte reduce; 1,292,544/9,400,320-byte per-rank gather inputs). +The trace kernel name ends in RING_LL even for these calls: that name alone +is not evidence of actual protocol. This explains why forcing Simple in the +previous handoff gave no benefit; no protocol override is promoted. + +The dominant grouped down kernel uses 255 registers/thread, 128 threads +and 32784 bytes shared memory per CTA. Register capacity bounds it to two +resident CTAs/SM, or eight warps (12.5% of 64); this is a resource ceiling, +not measured achieved occupancy. It motivates a bounded short-K tile screen +inside the existing provider, with smaller row tiles and the existing padded +activation loader. Its accelerated subset is grouped M>=512 routed rows, +N=2560 and K=160, for IQ4/group32 and affine U2/group32; other shapes retain +the existing tuner candidates. New GPU correctness and timing are pending. + +The first candidate's 260089-token quality prompt passes retrieval, arithmetic, +normal EOS and repetition checks (181 output tokens; maximum repeated ten-word +span 1). Hash: +`2c78b9e46e325dfd708b4d8c8f8c7c8fa2b9696bcb3d54f2a62313c70fb76665`. +The full response takes 821.64 s, under the existing 131072-token staging limit +and direct-host fallback beyond it. This is a boundary quality pass, not a +paired speed comparison; matching control is still pending. The owned API, +workers and Nsight launcher are stopped, with process ids archived in +`profile-stop-tree.json`. The remote trace files were removed only after +matching local SHA256 checks; local raw artifacts are `candidate-matched.*`. + +## Execution discipline + +Microbenchmarks acquire the existing GPU locks and start only on an empty GPU. +Other tasks' services are left to their owners. Build output, native tuning, +Triton, and Python compiler caches are task-owned. The shared venv is unchanged. +Candidate native code is built through the normal CMake `_C` target; no private +sidecar kernel or preload is an accepted artifact. + +The first native candidate installs `_C` SHA256 +`02dc9ff2fca7fb2c154448b5f21d537ee29a547b0ade6f5b5e4018759a56bb63`; +control is `5f0d7a83415532ec817c3e23a7b24b5015d7cbacafcd3d452fabbc28efeadd5b`. +The padded candidate installs `_C` SHA256 +`fc0c457f39c2bb7b788c614b1d67c8226922a34ec44aef6e20d005c3c7c0e7cf`. +The build reused three unchanged Marlin objects after verifying identical +normalized compiler commands, all 15 relocated source dependencies per +object, object hashes, and an up-to-date original Ninja dependency graph +(5864 dependencies/object). Those objects were copied into an owned cache; +an audited compiler launcher lets the normal CMake `_C` target consume them +and records the relocated dependencies. It invokes the real compiler on any +mismatch. All modified code was freshly compiled. This is an incremental +build, not a claim that every unchanged object was rebuilt from scratch. +The installed extension has no RPATH/RUNPATH or private DSO dependencies; +its dependencies are the declared CUDA/Torch libraries and system C++/C +runtime. Build-cache files are not runtime dependencies. + +The first isolated baseline startup compiled FlashQLA in an empty task cache. +When checkpoint loading overlapped four native compiler processes, the 5 GiB +host-memory guard stopped only the owned service tree. No benchmark request +completed in that attempt. Subsequent full-model startup must wait until the +large native build completes; low-memory microbenchmarks can run concurrently. +The effective unchanged residency log is 402/512 experts plus eight spare GPU +slots, 136 host slots, 9.46 GiB GPU expert slots and 3.14 GiB pinned host per +rank, matching the historical deployment with configured fraction 0.8. + +Draft PR: . The first candidate +passed matched 32K quality and unprofiled performance checks; 6K tokens/s is +not achieved. Paired 128K/near-256K qualification and the newly identified +segment-policy fix remain pending. The original deployment is unchanged. diff --git a/benchmarks/flashnext_prefill/analyze_mtp_verifier.py b/benchmarks/flashnext_prefill/analyze_mtp_verifier.py new file mode 100644 index 0000000000..d6f3c20e93 --- /dev/null +++ b/benchmarks/flashnext_prefill/analyze_mtp_verifier.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diagnostic MTP4 verifier phases, aligned by step ordinal across all TP4 ranks.""" + +import argparse +import json +import statistics +from pathlib import Path + +import regex as re + +p = argparse.ArgumentParser() +p.add_argument("log", type=Path) +p.add_argument("requests", type=Path) +p.add_argument("output", type=Path) +a = p.parse_args() + + +def stats(v): + x = sorted(v) + return dict( + n=len(x), + mean=statistics.mean(x), + p50=statistics.median(x), + p90=x[round(0.9 * (len(x) - 1))], + p99=x[round(0.99 * (len(x) - 1))], + max=x[-1], + ) + + +raw = a.log.read_bytes() +result = [] +for request in json.loads(a.requests.read_text()): + c = request["concurrency"] + groups = {} + for line in ( + raw[request["profile_byte_start"] : request["profile_byte_end"]] + .decode(errors="replace") + .splitlines() + ): + match = re.search( + r"Worker_TP(\d+).*SM70 V2 MTP profile interval_avg_ms " + r"calls=(\d+) interval_calls=1 tokens=(\d+) drafts=(\d+) (.*)", + line, + ) + if not match: + continue + rank, ordinal, tokens, drafts = map(int, match.groups()[:4]) + if (tokens, drafts) != (5 * c, 4 * c): + continue + values = {k: float(v) for k, v in re.findall(r"([a-z_]+)=([0-9.]+)", match[5])} + if rank in groups.setdefault(ordinal, {}): + raise ValueError("Duplicate rank ordinal") + groups[ordinal][rank] = values + complete = [(n, g) for n, g in sorted(groups.items()) if set(g) == {0, 1, 2, 3}] + if request["repetition"] < 0: + continue + if len(complete) < 16: + raise ValueError(("Too few full-concurrency verifier rounds", c, len(complete))) + records = [] + for ordinal, ranks in complete: + critical = max(ranks, key=lambda r: ranks[r]["target_verifier_gpu"]) + rec = dict( + ordinal=ordinal, critical_rank=critical, ranks=ranks, **ranks[critical] + ) + rec["rank_max_verifier_wall_ms"] = max( + r["target_verifier_wall_cpu"] for r in ranks.values() + ) + records.append(rec) + fields = ( + "target_verifier_gpu", + "rank_max_verifier_wall_ms", + "target_forward", + "target_sample", + "target_state_update", + "draft_total", + "total_gpu", + ) + row = dict( + concurrency=c, + repetition=request["repetition"], + complete_rounds=len(complete), + incomplete_rounds=len(groups) - len(complete), + phases={k: stats([r[k] for r in records]) for k in fields}, + per_round=records, + ) + result.append(row) + print(c, request["repetition"], json.dumps(row["phases"]), flush=True) +a.output.write_text( + json.dumps( + dict( + definition=( + "Per-ordinal TP4 maximum GPU event sum: target forward + sampling " + "+ state commit. Diagnostic events/fences are enabled. Draft is " + "separate; stage statistics are taken on the critical verifier rank." + ), + cases=result, + ), + indent=2, + ) +) diff --git a/benchmarks/flashnext_prefill/cold_dma.py b/benchmarks/flashnext_prefill/cold_dma.py new file mode 100644 index 0000000000..5bc35a8955 --- /dev/null +++ b/benchmarks/flashnext_prefill/cold_dma.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen copying only occupied cold slots, retaining their staging offsets. + +Uses three real TP4 layer-zero shard sizes and a synthetic fragmented slot +map. Timing includes Python/native submission of every contiguous range. +This is a DMA screen; model transcode overlap must be measured separately. +""" + +import argparse +import json +from functools import partial +from pathlib import Path + +import torch +from moe_tiles import artifact_identity, measure + +from vllm.model_executor.layers.quantization.gguf_expert_cache import _native +from vllm.transformers_utils.gguf_tensor_reader import quant_size + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + gate_block, gate_size = quant_size(18) + down_block, down_size = quant_size(20) + slot_bytes = [160 * (2560 // gate_block) * gate_size] * 2 + [ + 2560 * (160 // down_block) * down_size + ] + slots = 136 + host = [ + torch.empty(slots, size, dtype=torch.uint8).pin_memory() for size in slot_bytes + ] + for tensor in host: + tensor.random_(0, 256) + stage = [torch.empty_like(tensor, device="cuda") for tensor in host] + copy = _native().copy_h2d + stream = torch.cuda.current_stream().cuda_stream + report = {"artifact": artifact_identity(), "slot_bytes": slot_bytes, "cases": []} + for cold in (110, 104): + selected = torch.randperm(slots, generator=torch.Generator().manual_seed(9683))[ + :cold + ] + occupied = set(selected.tolist()) + extent = max(occupied) + 1 + ranges = [] + for slot in sorted(occupied): + if ranges and ranges[-1][1] == slot: + ranges[-1][1] += 1 + else: + ranges.append([slot, slot + 1]) + + def submit(spans): + for src, dst, size in zip(host, stage, slot_bytes): + for begin, end in spans: + copy( + dst.data_ptr() + begin * size, + src.data_ptr() + begin * size, + (end - begin) * size, + stream, + ) + + bulk = partial(submit, [(0, extent)]) + sparse = partial(submit, ranges) + for dst in stage: + dst.zero_() + sparse() + torch.accelerator.synchronize() + for src, dst in zip(host, stage): + actual = dst.cpu() + torch.testing.assert_close(actual[selected], src[selected], rtol=0, atol=0) + holes = torch.tensor(sorted(set(range(slots)) - occupied)) + assert (actual[holes] == 0).all() + row = { + "cold_slots": cold, + "extent": extent, + "ranges": len(ranges), + "bulk_bytes": extent * sum(slot_bytes), + "selected_bytes": cold * sum(slot_bytes), + "bulk": measure(bulk, 10), + "selected": measure(sparse, 10), + "bulk_repeat": measure(bulk, 10), + "cold_exact_holes_untouched": True, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/collectives.py b/benchmarks/flashnext_prefill/collectives.py new file mode 100644 index 0000000000..b538f66158 --- /dev/null +++ b/benchmarks/flashnext_prefill/collectives.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP4 prefill-sized collective screen; launch with torchrun on four idle GPUs. + +Uses the same PyNccl and custom-AR implementations as serving. Custom AR's larger +scratch allocation is experimental here, not a change to runtime admission. +""" + +import json +import os +import statistics +from functools import partial + +import torch +import torch.distributed as dist + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + + +def measure(fn): + samples = [] + for _ in range(5): + fn() + torch.accelerator.synchronize() + dist.barrier() + for _ in range(5): + begin, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + begin.record() + for _ in range(20): + fn() + end.record() + end.synchronize() + times = [None] * dist.get_world_size() + dist.all_gather_object(times, begin.elapsed_time(end) / 20) + samples.append(max(times)) + return {"median_ms": statistics.median(samples), "samples_ms": samples} + + +def main(): + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + nccl = PyNcclCommunicator(dist.group.WORLD, rank) + ca = CustomAllreduce(dist.group.WORLD, rank, max_size=48 * 1024**2) + assert not ca.disabled and not nccl.disabled + try: + for rows in (7344, 5347): + x = torch.randn(rows, 2560, device="cuda", dtype=torch.half) * 0.03 + control, candidate = torch.empty_like(x), torch.empty_like(x) + run_nccl = partial(nccl.all_reduce, x, control) + run_custom = partial(ca.all_reduce, x, out=candidate) + result = { + "tokens": rows, + "reduce_bytes": x.numel() * x.element_size(), + "nccl_all_reduce": measure(run_nccl), + "custom_all_reduce_eager": measure(run_custom), + } + difference = (candidate.float() - control.float()).norm() + result["relative_l2"] = (difference / control.float().norm()).item() + result["finite"] = bool(torch.isfinite(candidate).all()) + # Two HC packets per module: N=88 and N=640 per rank. + for n in (88, 640): + packet = torch.randn(rows, n, device="cuda", dtype=torch.half) + gathered = torch.empty(4 * rows, n, device="cuda", dtype=torch.half) + result[f"nccl_all_gather_n{n}"] = measure( + partial(nccl.all_gather, gathered, packet) + ) + if rank == 0: + print(json.dumps(result), flush=True) + finally: + torch.accelerator.synchronize() + ca.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/dense_blas.py b/benchmarks/flashnext_prefill/dense_blas.py new file mode 100644 index 0000000000..68e4cfd015 --- /dev/null +++ b/benchmarks/flashnext_prefill/dense_blas.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare existing fused and dequantize+cuBLAS dense providers on real TP4 weights. + +BLAS timing includes dequantization into its scratch on every call. This is a +provider screen, not a model qualification or a persistent FP16 weight cache. +""" + +import argparse +import json +from functools import partial +from pathlib import Path + +import numpy as np +import torch +from moe_tiles import artifact_identity, measure, relative_l2 + +from vllm.model_executor.layers.quantization.gguf_transcode import transcode_affine +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, dequantize + +CASES = [ + ("blk.0.attn_qkv.weight", 0), + ("blk.1.attn_qkv.weight", 0), + ("blk.0.ssm_out.weight", 1), + ("blk.1.attn_gate.weight", 0), + ("blk.5.ffn_up_shexp.weight", 0), +] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347]) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + reader = GGUFReader(args.model) + report = {"artifact": artifact_identity(), "cases": []} + for name, axis in CASES: + tensor = next(t for t in reader.tensors if t.name == name) + raw = np.ascontiguousarray(tensor.data) + kind = int(tensor.tensor_type) + canonical = transcode_affine(raw, kind).tp_slice(0, 4, axis=axis) + n, k = canonical.codes.shape + dense = dequantize(raw, kind) + dense = dense[:n] if axis == 0 else dense[:, :k] + reference_weight = torch.from_numpy(np.ascontiguousarray(dense)).cuda() + w, s, meta = torch.ops._C.gguf_affine_sm70_prepare( + torch.from_numpy(canonical.codes).cuda(), + torch.from_numpy(canonical.scales).cuda(), + torch.from_numpy(canonical.mins).cuda(), + canonical.bits, + canonical.group_size, + ) + k_ld, s_ld = meta.tolist() + scratch = torch.empty(k, n, dtype=torch.half, device="cuda") + for m in args.tokens: + torch.manual_seed(9683 + m) + x = (torch.randn(m, k, device="cuda") * 0.125).half() + fused_out = torch.empty(m, n, dtype=torch.half, device="cuda") + blas_out = torch.empty_like(fused_out) + + fused = partial( + torch.ops._C.gguf_affine_gemm_sm70_out, + fused_out, + x, + w, + s, + canonical.bits, + k_ld, + s_ld, + canonical.group_size, + ) + blas = partial( + torch.ops._C.gguf_affine_blas_sm70_out, + blas_out, + x, + w, + s, + canonical.bits, + scratch, + canonical.group_size, + ) + + fused() + blas() + expected = x[:16].float() @ reference_weight.float().T + errors = { + "fused_fp32_max_rel_l2": relative_l2(fused_out[:16], expected), + "blas_fp32_max_rel_l2": relative_l2(blas_out[:16], expected), + "paired_max_rel_l2": relative_l2(blas_out, fused_out), + } + passed = ( + bool(torch.isfinite(fused_out).all() and torch.isfinite(blas_out).all()) + and max(errors.values()) <= 0.003 + ) + row = { + "tensor": name, + "source_type": kind, + "tp_axis": axis, + "shape_mnk": [m, n, k], + "scratch_bytes": scratch.numel() * scratch.element_size(), + "fused": measure(fused, args.iterations), + "dequant_blas": measure(blas, args.iterations), + "fused_repeat": measure(fused, args.iterations), + **errors, + "passed": passed, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise SystemExit(0 if all(row["passed"] for row in report["cases"]) else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/dense_segments.py b/benchmarks/flashnext_prefill/dense_segments.py new file mode 100644 index 0000000000..471735bd3f --- /dev/null +++ b/benchmarks/flashnext_prefill/dense_segments.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure prefill after the serving segment-bank restoration, including DQ. + +Both arms keep small-M segment storage. The control omits only the admitted +BLAS bands, reproducing the missed serving route without another checkpoint. +""" + +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +from dense_blas import CASES +from moe_tiles import artifact_identity, measure, relative_l2 + +from tools.sm70.native_trace import NativeDispatchTrace +from vllm.config import set_current_vllm_config +from vllm.config.kernel import KernelConfig +from vllm.model_executor.layers.quantization.gguf_turbomind import ( + _prepared_gguf_mixed_projection, + prepare_gguf_projections, + prepared_projection_arguments, +) +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, dequantize + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + config = SimpleNamespace( + kernel_config=KernelConfig(), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + report = {"artifact": artifact_identity(), "cases": []} + reader = GGUFReader(args.model) + with set_current_vllm_config(config): + for name, axis in CASES[:4]: + tensor = next(t for t in reader.tensors if t.name == name) + raw = np.asarray(tensor.data) + raw = raw[: raw.shape[0] // 4] if axis == 0 else raw[:, : raw.shape[1] // 4] + raw = np.ascontiguousarray(raw) + source = int(tensor.tensor_type) + reference_weight = torch.from_numpy(dequantize(raw, source).copy()).cuda() + projection = prepare_gguf_projections( + [(torch.from_numpy(raw).cuda(), source)], torch.float16, True, 8 + )[0] + assert hasattr(projection, "segment_format") + actual = prepared_projection_arguments([projection]) + assert actual[3][0] == 3 and actual[5] == [2048, -1] + control = list(actual) + control[3] = actual[3][:-1] + [0] + control[5] = [] + k, n = projection.kernel.config.partition_weight_shape + for m in (7344, 5347): + torch.manual_seed(9683 + m) + x = (torch.randn(m, k, device="cuda") * 0.125).half() + + def fused(x=x, control=control): + return torch.ops.vllm.prepared_gguf_mixed_projection(x, *control) + + def blas(x=x, actual=actual): + return torch.ops.vllm.prepared_gguf_mixed_projection(x, *actual) + + old = fused() + with NativeDispatchTrace() as trace: + observed = _prepared_gguf_mixed_projection(x, *actual) + new = blas() + torch.testing.assert_close(new, observed, rtol=0, atol=0) + assert any( + r["operator"] == "_C::gguf_affine_blas_sm70_out" + and r["evidence"] == "cuda_dispatch" + for r in trace.records + ) + reference = x[:16].float() @ reference_weight.float().T + errors = { + "control_fp32_max_rel_l2": relative_l2(old[:16], reference), + "candidate_fp32_max_rel_l2": relative_l2(new[:16], reference), + "paired_max_rel_l2": relative_l2(new, old), + } + row = { + "tensor": name, + "source_type": source, + "shape_mnk": [m, n, k], + "restore_fused": measure(fused, 5), + "restore_blas": measure(blas, 5), + "restore_fused_repeat": measure(fused, 5), + **errors, + "passed": bool(torch.isfinite(new).all()) + and max(errors.values()) <= 0.003, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise SystemExit(0 if all(r["passed"] for r in report["cases"]) else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/hc_pipeline.py b/benchmarks/flashnext_prefill/hc_pipeline.py new file mode 100644 index 0000000000..7b841988ee --- /dev/null +++ b/benchmarks/flashnext_prefill/hc_pipeline.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen overlapping TP4 HC projections and gathers, without a model restart. + +Launch with torchrun on four idle V100s. This uses the production HC arithmetic +and NCCL provider with synthetic checkpoint-shaped inputs. It is a throughput +screen, not a model quality or runtime-state qualification. +""" + +import json +import os + +import torch +import torch.distributed as dist +from collectives import measure +from moe_tiles import relative_l2 + +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator +from vllm.models.qwen4_exp.nvidia.ops.hc import ( + hc_assemble_ranks, + hc_gate_mix_local, + hc_packet_silu, +) + +ORDER = (0, 1, 2, 3) + + +def sequential(nccl, x, down, up, rank): + local = torch.mm(x, down.T) + packet = local.new_empty((4 * x.shape[0], 88)) + nccl.all_gather(packet, local) + lora, injection = hc_packet_silu(packet, ORDER, 4) + local_gate = torch.mm(lora, up.T) + mixed = hc_gate_mix_local(x, local_gate, rank, 4) + gathered = mixed.new_empty((4 * x.shape[0], 640)) + nccl.all_gather(gathered, mixed) + return hc_assemble_ranks(gathered, ORDER), injection + + +class Pipeline: + def __init__(self, nccl, x, down, up, rank, chunks, overlap=True): + self.nccl, self.down, self.up, self.rank = nccl, down, up, rank + self.x = x.chunk(chunks) + self.comm = torch.cuda.Stream() if overlap else torch.cuda.current_stream() + self.ready = [torch.cuda.Event() for _ in self.x] + self.packet_done = [torch.cuda.Event() for _ in self.x] + self.mixed_ready = [torch.cuda.Event() for _ in self.x] + self.gather_done = [torch.cuda.Event() for _ in self.x] + self.packet = [v.new_empty((4 * v.shape[0], 88)) for v in self.x] + self.gathered = [v.new_empty((4 * v.shape[0], 640)) for v in self.x] + + def __call__(self): + main = torch.cuda.current_stream() + # Retain the inputs until all their asynchronous gathers are queued. + locals_, mixed, injections = [], [], [] + for i, x in enumerate(self.x): + locals_.append(torch.mm(x, self.down.T)) + self.ready[i].record(main) + with torch.cuda.stream(self.comm): + self.comm.wait_event(self.ready[i]) + self.nccl.all_gather(self.packet[i], locals_[-1]) + locals_[-1].record_stream(self.comm) + self.packet_done[i].record(self.comm) + for i, x in enumerate(self.x): + main.wait_event(self.packet_done[i]) + lora, injection = hc_packet_silu(self.packet[i], ORDER, 4) + injections.append(injection) + gate = torch.mm(lora, self.up.T) + mixed.append(hc_gate_mix_local(x, gate, self.rank, 4)) + self.mixed_ready[i].record(main) + with torch.cuda.stream(self.comm): + self.comm.wait_event(self.mixed_ready[i]) + self.nccl.all_gather(self.gathered[i], mixed[-1]) + mixed[-1].record_stream(self.comm) + self.gather_done[i].record(self.comm) + outputs = [] + for i, gathered in enumerate(self.gathered): + main.wait_event(self.gather_done[i]) + outputs.append(hc_assemble_ranks(gathered, ORDER)) + return torch.cat(outputs), torch.cat(injections) + + +def main(): + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + nccl = PyNcclCommunicator(dist.group.WORLD, rank) + assert not nccl.disabled + passed = True + try: + for rows in (7344, 5347): + torch.manual_seed(rows) + x = torch.randn(rows, 10240, device="cuda", dtype=torch.half) * 0.3 + torch.manual_seed(100 + rank) + down = torch.randn(88, 10240, device="cuda", dtype=torch.half) * 0.01 + up = torch.randn(2560, 320, device="cuda", dtype=torch.half) * 0.01 + + def control(x=x, down=down, up=up): + return sequential(nccl, x, down, up, rank) + + expected = control() + result = {"tokens": rows, "sequential": measure(control)} + for chunks in (2, 4): + candidate = Pipeline(nccl, x, down, up, rank, chunks) + actual = candidate() + same_tiles = Pipeline(nccl, x, down, up, rank, chunks, overlap=False)() + error = [relative_l2(a, b) for a, b in zip(actual, expected)] + serial_error = max( + relative_l2(a, b) for a, b in zip(actual, same_tiles) + ) + finite = all(bool(torch.isfinite(a).all()) for a in actual) + errors = [None] * 4 + dist.all_gather_object(errors, (error, serial_error, finite)) + numerical_pass = all( + max(e) <= 0.001 and s == 0 and f for e, s, f in errors + ) + passed &= numerical_pass + # Failed numerical screens remain failed. Timing them decides + # whether further diagnosis is worthwhile, not admission. + result[f"chunks_{chunks}"] = measure(candidate) | { + "max_rank_output_relative_l2": max(e[0] for e, _, _ in errors), + "max_rank_injection_relative_l2": max(e[1] for e, _, _ in errors), + "overlap_vs_serial_same_tiles": max(s for _, s, _ in errors), + "finite_all_ranks": all(f for _, _, f in errors), + "numerical_pass": numerical_pass, + } + if rank == 0: + print(json.dumps(result), flush=True) + finally: + torch.accelerator.synchronize() + dist.destroy_process_group() + if not passed: + raise SystemExit("Numerical screen failed; candidate is not qualified") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/hc_projection.py b/benchmarks/flashnext_prefill/hc_projection.py new file mode 100644 index 0000000000..df3c412aa5 --- /dev/null +++ b/benchmarks/flashnext_prefill/hc_projection.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test whether HC's narrow N=88 projection benefits from more K partitions. + +Synthetic FP16 operands, serving dimensions; all packing/reduction costs are +included. This is a rejection/admission experiment, not model qualification. +""" + +import argparse +import json +from functools import partial + +import torch +from moe_tiles import measure, relative_l2 + + +def partitioned(x, weight, splits): + rows, k = x.shape + n = weight.shape[0] + partials = torch.bmm( + x.view(rows, splits, k // splits).permute(1, 0, 2), + weight.view(n, splits, k // splits).permute(1, 2, 0), + out_dtype=torch.float32, + ) + return partials.sum(dim=0).half() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347]) + args = parser.parse_args() + torch.backends.cuda.matmul.allow_tf32 = False + for rows in args.tokens: + gen = torch.Generator(device="cuda").manual_seed(rows) + x = torch.randn(rows, 10240, device="cuda", dtype=torch.half, generator=gen) + weight = ( + torch.randn(88, 10240, device="cuda", generator=gen) / 10240**0.5 + ).half() + expected = x.float() @ weight.float().T + baseline = torch.mm(x, weight.T) + result = {"tokens": rows, "candidates": []} + result["baseline"] = measure(partial(torch.mm, x, weight.T), 20) + result["baseline_rel_l2"] = relative_l2(baseline, expected) + for splits in (2, 4, 8): + candidate = partitioned(x, weight, splits) + result["candidates"].append( + { + "splits": splits, + **measure(partial(partitioned, x, weight, splits), 20), + "fp32_rel_l2": relative_l2(candidate, expected), + "baseline_rel_l2": relative_l2(candidate, baseline), + "finite": bool(torch.isfinite(candidate).all()), + } + ) + print(json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/measure_mtp_concurrency.py b/benchmarks/flashnext_prefill/measure_mtp_concurrency.py new file mode 100644 index 0000000000..a1b5f249af --- /dev/null +++ b/benchmarks/flashnext_prefill/measure_mtp_concurrency.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Matched concurrent MTP requests, with an explicit common decode interval. + +Token timestamps are HTTP arrival times, not GPU verifier timers. Warm the same +prompts first; use the worker step profiler separately for phase attribution. +""" + +import argparse +import hashlib +import json +import os +import statistics +import threading +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import regex as re + + +def request(base, path, body=None): + data = json.dumps(body).encode() if body is not None else None + headers = {"Content-Type": "application/json"} + if key := os.environ.get("VLLM_API_KEY"): + headers["Authorization"] = "Bearer " + key + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + return opener.open(urllib.request.Request(base + path, data, headers), timeout=900) + + +def counters(base): + with request(base, "/metrics") as response: + text = response.read().decode() + return { + name: sum( + float(value) + for value in re.findall( + rf"^vllm:spec_decode_num_{name}_total\{{[^}}]*\}} ([0-9.e+]+)$", + text, + re.M, + ) + ) + for name in ("drafts", "draft_tokens", "accepted_tokens") + } + + +def run_one(base, case, output_tokens, barrier, seed): + body = { + "model": "flash-next-abliterated", + "messages": [{"role": "user", "content": case["prompt"]}], + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "seed": seed, + "max_tokens": output_tokens, + "chat_template_kwargs": {"enable_thinking": True}, + "stream": True, + "stream_options": {"include_usage": True}, + "return_token_ids": True, + } + barrier.wait() + started = time.perf_counter() + events = [] + content = [] + reasoning = [] + usage = None + finish = None + done = False + with request(base, "/v1/chat/completions", body) as response: + for raw in response: + line = raw.decode().strip() + if not line.startswith("data:"): + continue + if line == "data: [DONE]": + done = True + continue + obj = json.loads(line[5:]) + if obj.get("usage"): + usage = obj["usage"] + for choice in obj.get("choices", []): + finish = choice.get("finish_reason") or finish + delta = choice.get("delta", {}) + ids = choice.get("token_ids") or delta.get("token_ids") or [] + if ids: + events.append({"at": time.perf_counter(), "token_ids": ids}) + content.append(delta.get("content") or "") + reasoning.append(delta.get("reasoning") or "") + ended = time.perf_counter() + if not done or usage is None or not events: + raise RuntimeError("Incomplete stream or missing emitted token IDs") + emitted = sum(len(event["token_ids"]) for event in events) + if emitted != usage["completion_tokens"]: + raise RuntimeError(f"Token accounting differs: {emitted} versus {usage}") + first, last = events[0]["at"], events[-1]["at"] + decode_tokens = emitted - len(events[0]["token_ids"]) + return { + "seed": seed, + "prompt_sha256": hashlib.sha256(case["prompt"].encode()).hexdigest(), + "usage": usage, + "finish_reason": finish, + "ttft_s": first - started, + "total_s": ended - started, + "decode_s": last - first, + "decode_tokens": decode_tokens, + "output_token_ms": 1000 * (last - first) / max(decode_tokens, 1), + "events": events, + "content": "".join(content), + "reasoning": "".join(reasoning), + } + + +def run_batch(base, case, concurrency, output_tokens, repetition, independent_prefix): + before = counters(base) + barrier = threading.Barrier(concurrency) + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=concurrency) as pool: + futures = [] + for i in range(concurrency): + worker_case = dict(case) + if independent_prefix: + worker_case["prompt"] = f"[Independent request {i}]\n" + case["prompt"] + futures.append( + pool.submit( + run_one, base, worker_case, output_tokens, barrier, 20261014 + i + ) + ) + rows = [future.result() for future in futures] + time.sleep(0.2) # Let the final asynchronous scheduler counters settle. + after = counters(base) + # Exclude every request's first and last emission. All requests must still + # be decoding in the shared interval; this avoids calling tail drain C8. + common_start = max(row["events"][0]["at"] for row in rows) + common_end = min(row["events"][-1]["at"] for row in rows) + seconds = common_end - common_start + common_tokens = sum( + len(event["token_ids"]) + for row in rows + for event in row["events"] + if common_start < event["at"] < common_end + ) + valid_overlap = seconds > 0 and common_tokens >= concurrency * 16 + if repetition >= 0 and not valid_overlap: + raise RuntimeError( + "Insufficient common decode interval; increase output length" + ) + delta = {key: after[key] - before[key] for key in before} + return { + "concurrency": concurrency, + "prefix_mode": "independent" if independent_prefix else "shared", + "repetition": repetition, + "batch_elapsed_s": time.perf_counter() - start, + "mean_ttft_s": statistics.mean(row["ttft_s"] for row in rows), + "mean_output_token_ms": statistics.mean(row["output_token_ms"] for row in rows), + "common_decode_s": seconds if valid_overlap else None, + "common_decode_tokens": common_tokens, + "common_aggregate_tokens_per_s": common_tokens / seconds + if valid_overlap + else None, + "common_per_request_token_ms": ( + 1000 * seconds * concurrency / common_tokens if valid_overlap else None + ), + "accepted_length": 1 + delta["accepted_tokens"] / max(delta["drafts"], 1), + "mtp_metrics_delta": delta, + "requests": rows, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18041") + parser.add_argument("--cases", type=Path, required=True) + parser.add_argument("--case", default="i32768") + parser.add_argument("--concurrency", nargs="+", type=int, default=[1, 2, 4, 6, 8]) + parser.add_argument("--output-tokens", type=int, default=384) + parser.add_argument("--repetitions", type=int, default=3) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--independent-prefix", action="store_true") + parser.add_argument("--service-log", type=Path) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + case = next( + row for row in json.loads(args.cases.read_text()) if row["name"] == args.case + ) + results = [] + for concurrency in args.concurrency: + for repetition in range(-args.warmups, args.repetitions): + log_start = args.service_log.stat().st_size if args.service_log else None + row = run_batch( + args.base_url, + case, + concurrency, + args.output_tokens, + repetition, + args.independent_prefix, + ) + results.append(row) + if args.service_log: + row.update( + profile_byte_start=log_start, + profile_byte_end=args.service_log.stat().st_size, + ) + args.out.write_text(json.dumps(results, indent=2)) + print( + json.dumps({k: v for k, v in row.items() if k != "requests"}), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/measure_prefill.py b/benchmarks/flashnext_prefill/measure_prefill.py new file mode 100644 index 0000000000..fbebad9d23 --- /dev/null +++ b/benchmarks/flashnext_prefill/measure_prefill.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Matched, cold-prefix prefill requests against an exclusively owned server. + +Use the same --seed/--words/--warmup/--trials for control and candidate. Word +count is not token count: throughput uses the server's reported prompt_tokens. +--profile wraps only measured requests in the server's profiler window; do not +use profiled timing as an unprofiled end-to-end result. +""" + +import argparse +import hashlib +import json +import random +import statistics +import time +import urllib.request +from pathlib import Path + +WORDS = "alpha river stone cloud engine paper signal forest copper lantern orbit meadow" +METRICS = ( + "vllm:request_prefill_time_seconds", + "vllm:time_to_first_token_seconds", + "vllm:request_queue_time_seconds", +) + + +def request(base, path, body=None): + data = None if body is None else json.dumps(body).encode() + req = urllib.request.Request( + base + path, data, {"Content-Type": "application/json"} + ) + return urllib.request.urlopen(req, timeout=1800) + + +def metrics(base): + with request(base, "/metrics") as response: + lines = response.read().decode().splitlines() + result = {} + for line in lines: + if not line or line.startswith("#"): + continue + name = line.split("{", 1)[0].split(" ", 1)[0] + if name in {m + suffix for m in METRICS for suffix in ("_sum", "_count")}: + result[name] = result.get(name, 0) + float(line.rsplit(" ", 1)[1]) + return result + + +def metric_deltas(base, before): + deadline = time.monotonic() + 8 + while True: + after = metrics(base) + counts = { + m: after.get(m + "_count", 0) - before.get(m + "_count", 0) for m in METRICS + } + if all(n >= 1 for n in counts.values()) or time.monotonic() >= deadline: + break + time.sleep(0.2) + # More than one request means this measurement cannot be attributed. + return { + m.removeprefix("vllm:"): { + "count": counts[m], + "seconds": ( + after.get(m + "_sum", 0) - before.get(m + "_sum", 0) + if counts[m] == 1 + else None + ), + } + for m in METRICS + } + + +def run_one(args, index): + rng = random.Random(args.seed + index) + # An early request-specific nonce also avoids sharing warmup prefixes. + prompt = f"[{rng.randrange(10**9)}] " + " ".join( + rng.choice(WORDS.split()) for _ in range(args.words) + ) + body = { + "model": args.model, + "prompt": prompt, + "temperature": 0, + "max_tokens": 1, + "stream": True, + "stream_options": {"include_usage": True}, + } + # This endpoint is intended only for a server owned by this benchmark. + with request(args.base_url, "/reset_prefix_cache", {}) as response: + response.read() + before = metrics(args.base_url) + start = time.perf_counter() + first_choice = first_text = None + usage, choices = None, [] + with request(args.base_url, "/v1/completions", body) as response: + for line in response: + if not line.startswith(b"data: "): + continue + data = line[6:].strip() + if data == b"[DONE]": + break + event = json.loads(data) + if "error" in event: + raise RuntimeError(event["error"]) + if event.get("choices"): + choices.extend(event["choices"]) + elapsed = time.perf_counter() - start + if first_choice is None: + first_choice = elapsed + if first_text is None and any(c.get("text") for c in event["choices"]): + first_text = elapsed + if event.get("usage"): + usage = event["usage"] + total = time.perf_counter() - start + if first_choice is None or usage is None: + raise RuntimeError("Missing completion choice or final usage") + return { + "index": index, + "prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(), + "usage": usage, + "first_choice_seconds": first_choice, + "first_text_seconds": first_text, + "response_seconds": total, + "prompt_tokens_per_first_choice_second": usage["prompt_tokens"] / first_choice, + "choices": choices, + "server_metrics": metric_deltas(args.base_url, before), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8002") + parser.add_argument("--model", default="flash-next") + parser.add_argument("--words", type=int, default=32000) + parser.add_argument("--seed", type=int, default=9683) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--profile", action="store_true") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.words < 1 or args.warmup < 0 or args.trials < 1: + parser.error("words/trials must be positive and warmup nonnegative") + report = {"config": vars(args) | {"output": str(args.output)}, "requests": []} + for i in range(args.warmup + args.trials): + profiling = args.profile and i >= args.warmup + if profiling: + with request(args.base_url, "/start_profile", {}) as response: + response.read() + try: + result = run_one(args, i) + finally: + if profiling: + with request(args.base_url, "/stop_profile", {}) as response: + response.read() + result["warmup"] = i < args.warmup + report["requests"].append(result) + print(json.dumps(result), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + measured = [r for r in report["requests"] if not r["warmup"]] + report["summary"] = { + "profiled": args.profile, + "median_first_choice_seconds": statistics.median( + r["first_choice_seconds"] for r in measured + ), + "median_prompt_tokens_per_first_choice_second": statistics.median( + r["prompt_tokens_per_first_choice_second"] for r in measured + ), + } + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report["summary"]), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_blas.py b/benchmarks/flashnext_prefill/moe_blas.py new file mode 100644 index 0000000000..09b12d221e --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_blas.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen bounded expert-batch FP16 GEMMs against the quantized MoE chain. + +Real TP4 weights and routes; both arms exclude routing and host DMA. The BLAS +arm includes dequantization, padded input gather, scatter and weighted unroute. +Its contiguous raw GPU banks are an optimistic source layout: production's +resident/cold slots would need additional gathering. This is not a new backend. +""" + +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +import torch +from moe_tiles import ( + EXPERTS, + HIDDEN, + INTERMEDIATE, + TOPK, + artifact_identity, + load_layer, + measure, + relative_l2, +) + +from vllm.model_executor.layers.quantization.gguf import dequantize_gguf_rows +from vllm.model_executor.layers.quantization.gguf_expert_cache import ( + TurboMindPrefill, + _unroute, +) +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, quant_size +from vllm.triton_utils import tl, triton + + +@triton.jit +def _pad_gather( + x, + rows, + offsets, + padded, + FIRST: tl.constexpr, + PAD: tl.constexpr, + K: tl.constexpr, + BLOCK: tl.constexpr, +): + expert = tl.program_id(0) + pos = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + r, col = pos // K, pos % K + begin = tl.load(offsets + FIRST + expert) + end = tl.load(offsets + FIRST + expert + 1) + token = tl.load(rows + begin + r, (r < PAD) & (begin + r < end), 0) + value = tl.load(x + token * K + col, (r < PAD) & (begin + r < end), 0) + tl.store(padded + expert * PAD * K + pos, value, r < PAD) + + +@triton.jit +def _scatter( + padded, + offsets, + out, + FIRST: tl.constexpr, + PAD: tl.constexpr, + K: tl.constexpr, + BLOCK: tl.constexpr, +): + expert = tl.program_id(0) + pos = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + r, col = pos // K, pos % K + begin = tl.load(offsets + FIRST + expert) + end = tl.load(offsets + FIRST + expert + 1) + valid = (r < PAD) & (begin + r < end) + value = tl.load(padded + expert * PAD * K + pos, valid, 0) + tl.store(out + (begin + r) * K + col, value, valid) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--layer", type=int, default=0) + parser.add_argument("--tokens", type=int, default=7344) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + info, tables, banks, raw, left = load_layer(GGUFReader(args.model), args.layer) + runner = TurboMindPrefill.__new__(TurboMindPrefill) + runner.cache = SimpleNamespace(experts=EXPERTS, prefill_tokens=4096) + runner.expert_ids = torch.arange(EXPERTS + 1, dtype=torch.int32, device="cuda") + runner.tuned = set() + runner._tune( + args.layer, + info, + tables, + torch.zeros(8192, HIDDEN, device="cuda", dtype=torch.half), + ) + torch.manual_seed(args.tokens) + x = (torch.randn(args.tokens, HIDDEN, device="cuda") * 0.5).half() + scores, ids = torch.randn(args.tokens, EXPERTS, device="cuda").topk(TOPK, dim=1) + weights = scores.softmax(dim=1) + route = runner._route(ids, 0, EXPERTS) + offsets, _, rows, _, inverse, _ = route + pad = triton.cdiv(int((offsets[1:] - offsets[:-1]).max()), 32) * 32 + subroutes = [ + ( + t, + min(t + 4096, args.tokens), + runner._subroute(route, t, min(t + 4096, args.tokens)), + ) + for t in range(0, args.tokens, 4096) + ] + control = torch.empty_like(x) + + def baseline(): + hidden = runner._gate_up(info, tables, x, route, False) + for first, last, sub in subroutes: + runner._down( + info, + tables, + hidden, + sub, + weights[first:last], + control[first:last], + False, + ) + + baseline() + raw_gu = torch.cat([torch.from_numpy(raw[p]).cuda() for p in ("gate", "up")], 1) + raw_down = torch.from_numpy(raw["down"]).cuda() + block, size = quant_size(info["down"][0]) + down_k = raw_down.shape[-1] // size * block + routed = torch.empty(args.tokens * TOPK, HIDDEN, device="cuda", dtype=torch.half) + output = torch.empty_like(x) + report = { + "artifact": artifact_identity(), + "layer": args.layer, + "tokens": args.tokens, + "padded_rows_per_expert": pad, + "control": measure(baseline, 3), + "cases": [], + } + for count in (32, 64, 128): + padded = torch.empty(count, pad, HIDDEN, device="cuda", dtype=torch.half) + + def candidate(count=count, padded=padded): + for first in range(0, EXPERTS, count): + gu = dequantize_gguf_rows( + raw_gu[first : first + count].flatten(0, 1), + info["gate"][0], + HIDDEN, + torch.half, + ).view(count, 2 * INTERMEDIATE, HIDDEN) + dn = dequantize_gguf_rows( + raw_down[first : first + count].flatten(0, 1), + info["down"][0], + down_k, + torch.half, + ).view(count, HIDDEN, down_k) + dn = dn[:, :, left : left + INTERMEDIATE] + grid = (count, triton.cdiv(pad * HIDDEN, 2048)) + _pad_gather[grid](x, rows, offsets, padded, first, pad, HIDDEN, 2048) + projected = torch.bmm( + padded, gu.transpose(1, 2), out_dtype=torch.float32 + ) + hidden = ( + torch.nn.functional.silu(projected[:, :, :INTERMEDIATE]) + * projected[:, :, INTERMEDIATE:] + ).half() + down = torch.bmm(hidden, dn.transpose(1, 2)) + _scatter[grid](down, offsets, routed, first, pad, HIDDEN, 2048) + _unroute[(args.tokens, triton.cdiv(HIDDEN, 1024))]( + routed, weights, inverse, output, HIDDEN, TOPK=TOPK, BLOCK=1024 + ) + + candidate() + error = relative_l2(output, control) + finite = bool(torch.isfinite(output).all()) + row = { + "experts_per_batch": count, + **measure(candidate, 3), + "control_max_rel_l2": error, + "finite": finite, + "passed": finite and error <= 0.003, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + report["control_repeat"] = measure(baseline, 3) + # Localize a rejected whole-chain screen before considering native work. + count = 64 + packed_gu = raw_gu[:count].flatten(0, 1) + packed_dn = raw_down[:count].flatten(0, 1) + + def dequant_gu(): + return dequantize_gguf_rows( + packed_gu, info["gate"][0], HIDDEN, torch.half + ).view(count, 2 * INTERMEDIATE, HIDDEN) + + def dequant_dn(): + return dequantize_gguf_rows( + packed_dn, info["down"][0], down_k, torch.half + ).view(count, HIDDEN, down_k)[:, :, left : left + INTERMEDIATE] + + gu, dn = dequant_gu(), dequant_dn() + padded = torch.empty(count, pad, HIDDEN, device="cuda", dtype=torch.half) + grid = (count, triton.cdiv(pad * HIDDEN, 2048)) + + def gather(): + _pad_gather[grid](x, rows, offsets, padded, 0, pad, HIDDEN, 2048) + + gather() + projected = torch.bmm(padded, gu.transpose(1, 2), out_dtype=torch.float32) + hidden = ( + torch.nn.functional.silu(projected[:, :, :INTERMEDIATE]) + * projected[:, :, INTERMEDIATE:] + ).half() + down = torch.bmm(hidden, dn.transpose(1, 2)) + parts = { + "dequant_gate_up": dequant_gu, + "dequant_down": dequant_dn, + "gather": gather, + "gate_up_bmm_fp32_output": lambda: torch.bmm( + padded, gu.transpose(1, 2), out_dtype=torch.float32 + ), + "gate_up_bmm_fp16_output": lambda: torch.bmm(padded, gu.transpose(1, 2)), + "down_bmm": lambda: torch.bmm(hidden, dn.transpose(1, 2)), + "scatter": lambda: _scatter[grid](down, offsets, routed, 0, pad, HIDDEN, 2048), + } + report["first_64_expert_components"] = { + name: measure(fn, 20) for name, fn in parts.items() + } + print(json.dumps(report["first_64_expert_components"]), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise SystemExit(0 if all(row["passed"] for row in report["cases"]) else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_dma.py b/benchmarks/flashnext_prefill/moe_dma.py new file mode 100644 index 0000000000..812cd20356 --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_dma.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare full-extent and occupied-slot DMA in the actual prefill pipeline. + +The eight-layer cache, residency, staging allocation, transcode, and GEMMs are +identical in both modes. Initial and fragmented post-decode maps are timed. +""" + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from moe_streams import build_cache +from moe_tiles import EXPERTS, HIDDEN, TOPK, artifact_identity, measure + +from vllm.model_executor.layers.quantization.gguf_expert_cache import SHARDS, _native +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader + + +class LaneProbe: + def __init__(self, cache): + self.cache, self.lane = cache, cache.lane + self.mode = "bulk" + + def __getattr__(self, name): + return getattr(self.lane, name) + + def stage_cold(self, i, shard, dst, stream): + if self.mode == "sparse": + return self.lane.stage_cold(i, shard, dst, stream) + cache = self.cache + layer = cache.layers[i] + rows = cache.tm.extent[i % cache.tm.VIEWS] + if rows: + _native().copy_h2d( + dst, + cache.host[layer, SHARDS[shard]].data_ptr(), + rows * cache.sizes[layer, SHARDS[shard]], + stream, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.manual_seed(9683) + cache = build_cache(GGUFReader(args.model)) + cache.lane = LaneProbe(cache) + x = (torch.randn(7344, HIDDEN, device="cuda") * 0.5).half() + scores, ids = torch.randn(7344, EXPERTS, device="cuda").topk(TOPK, dim=1) + weights = scores.softmax(dim=1) + native = Path(_native().__file__) + report = { + "artifact": artifact_identity(), + "cache_native_sha256": hashlib.sha256(native.read_bytes()).hexdigest(), + "cases": [], + "maps": {}, + } + + def forward(gap=0, count=8): + outputs = [] + for layer in cache.layers[:count]: + outputs.append(cache.prefill_moe(layer, x, ids, weights)) + if gap: + torch.cuda._sleep(gap) + return outputs + + def check(): + for actual, expected in zip(forward(), reference): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + assert cache.tm.error.item() == 0 + + try: + reference = [out.cpu() for out in forward()] + for phase in ("initial", "after_decode"): + if phase == "after_decode": + for _ in range(10): + for layer in cache.layers: + cache.resolve( + layer, torch.randperm(EXPERTS, device="cuda")[:50] + ) + torch.accelerator.synchronize() + report["maps"][phase] = [] + for i in range(len(cache.layers)): + host = cache.host_slot_of[i][cache.slot_of[i] < 0] + report["maps"][phase].append( + {"cold": host.numel(), "extent": int(host.max()) + 1} + ) + for gap in (0, 20000000): + for mode in ("bulk", "sparse", "bulk"): + torch.accelerator.synchronize() + cache.lane.mode = mode + cache.tm.prefetched.clear() + cache.tm.start() + check() + forward(gap, count=4) + check() # Restart while the interrupted pass has pending copies. + row = { + "phase": phase, + "mode": mode, + "gap_cycles": gap, + **measure(lambda gap=gap: forward(gap), 2), + "exact_output_and_restart": True, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + report["cache_stats"] = cache.stats() + assert report["cache_stats"]["errors"] == 0 + args.output.write_text(json.dumps(report, indent=2) + "\n") + finally: + torch.accelerator.synchronize() + cache.close() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_streams.py b/benchmarks/flashnext_prefill/moe_streams.py new file mode 100644 index 0000000000..48e27a62ca --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_streams.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen prefill stream dependencies with real weights and the native cache. + +Uses eight logical layers (four real TP4 rank-zero layers repeated) to wrap +the residency-view ring. All modes share kernels, weights, expert residency, +and bank/staging capacity. Optional GPU sleep models a between-layer delay, +not concurrent model compute. This is not an end-to-end throughput result. +""" + +import argparse +import json +import statistics +from pathlib import Path + +import torch +from moe_tiles import EXPERTS, HIDDEN, INTERMEDIATE, TOPK, artifact_identity, read_layer + +from vllm.model_executor.layers.quantization.gguf_expert_cache import ( + SHARDS, + ExpertCache, + TurboMindPrefill, +) +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader + + +class StreamProbe(TurboMindPrefill): + """Benchmark-only policies; production _issue/_locate/compute are reused.""" + + mode = "shared" + + def start(self): + super().start() + device = self.cache.device + self.cold_by_key = { + "gu": self.cold_stream, + "dn": torch.cuda.Stream(device, priority=-1), + } + self.resident_by_key = { + "gu": self.resident_stream, + "dn": torch.cuda.Stream(device, priority=-1), + } + self.located = [torch.cuda.Event() for _ in range(self.VIEWS)] + self.view_done = [torch.cuda.Event() for _ in range(self.VIEWS)] + for event in self.view_done: + event.record() + + def _locate(self, i): + p = i % self.VIEWS + if self.mode == "split": + # Down may still use this view after an interrupted pass. + self.resident_stream.wait_event(self.view_done[p]) + super()._locate(i) + self.located[p].record(self.resident_stream) + + def _issue(self, key, i): + self.cold_stream = self.cold_by_key["gu" if self.mode == "shared" else key] + self.resident_stream = self.resident_by_key[ + key if self.mode == "split" else "gu" + ] + if key == "dn" and self.mode == "split": + self.resident_stream.wait_event(self.located[i % self.VIEWS]) + super()._issue(key, i) + if key == "dn": + self.view_done[i % self.VIEWS].record(self.resident_stream) + + +def build_cache(reader): + real_layers = [0, 1, 17, 47] + sources = [read_layer(reader, layer) for layer in real_layers] + banks = [] + for layer in range(8): + raw, gate_type, down_type, left = sources[layer % len(sources)] + for shard, part in zip(SHARDS, ("gate", "up", "down")): + source = ( + (down_type, INTERMEDIATE, left) + if part == "down" + else (gate_type, HIDDEN, 0) + ) + banks.append((layer, shard, raw[part].shape, 0, source)) + cache = ExpertCache(banks, EXPERTS, 0.8, 8, 0.05, None, "cuda:0", "turbomind", 4096) + for layer in cache.layers: + raw = sources[layer % len(sources)][0] + for shard, part in zip(SHARDS, ("gate", "up", "down")): + cache.store_bank(layer, shard, torch.from_numpy(raw[part])) + cache.tm.__class__ = StreamProbe + cache.start() + return cache + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--tokens", type=int, default=7344) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--gap-cycles", type=int, nargs="+", default=[0, 20000000]) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.tokens < 1 or args.iterations < 1 or min(args.gap_cycles) < 0: + parser.error("invalid tokens, iterations, or gap cycles") + torch.manual_seed(9683) + cache = build_cache(GGUFReader(args.model)) + gen = torch.Generator(device="cuda").manual_seed(9683) + x = (torch.randn(args.tokens, HIDDEN, device="cuda", generator=gen) * 0.5).half() + scores, ids = torch.randn(args.tokens, EXPERTS, device="cuda", generator=gen).topk( + TOPK, dim=1 + ) + weights = scores.softmax(dim=1) + report = { + "artifact": artifact_identity(), + "tokens": args.tokens, + "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(), + "resident_experts": cache.resident0, + "host_slots": cache.host_slots, + "canonical_and_staging_bytes": cache.tm.storage_bytes(), + "cases": [], + } + + def forward(gap, count=8): + outputs = [] + for layer in cache.layers[:count]: + outputs.append(cache.prefill_moe(layer, x, ids, weights)) + if gap: + torch.cuda._sleep(gap) + return outputs + + def select(mode): + torch.accelerator.synchronize() + cache.tm.prefetched.clear() + cache.tm.mode = mode + cache.tm.start() + + def compare(outputs, reference): + # Stream changes preserve arithmetic: exact output is the contract. + for actual, expected in zip(outputs, reference): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + assert cache.tm.error.item() == 0 + + try: + # Warm native tuning, Triton, allocation pools and every source format. + reference = [out.cpu() for out in forward(0)] + for gap in args.gap_cycles: + for mode in ("shared", "cold_split", "split", "shared"): + select(mode) + compare(forward(gap), reference) + # Restart without draining side-stream work from this pass. + forward(gap, count=4) + compare(forward(gap), reference) + samples = [] + for _ in range(5): + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + for _ in range(args.iterations): + forward(gap) + end.record() + end.synchronize() + samples.append(begin.elapsed_time(end) / args.iterations) + case = { + "mode": mode, + "gap_cycles": gap, + "median_ms": statistics.median(samples), + "samples_ms": samples, + "exact_output_and_restart": True, + } + report["cases"].append(case) + print(json.dumps(case), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + # Exercise real decode misses, host write-backs, and a new snapshot. + for layer in cache.layers: + cache.resolve(layer, torch.arange(480, 512, device="cuda")) + torch.accelerator.synchronize() + for mode in ("shared", "cold_split", "split"): + select(mode) + compare(forward(0), reference) + report["after_decode_residency_change_exact"] = True + report["cache_stats"] = cache.stats() + assert report["cache_stats"]["errors"] == 0 + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report["cache_stats"]), flush=True) + finally: + torch.accelerator.synchronize() + cache.close() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_tiles.py b/benchmarks/flashnext_prefill/moe_tiles.py new file mode 100644 index 0000000000..b13fd544d2 --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_tiles.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real-weight TP4 MoE microbenchmark, using the serving route/subroute code. + +Run control/candidate in fresh processes with separate native tuning caches. +This isolates resident canonical-bank compute: it does not measure host DMA, +transcoding overlap, TP communication, or end-to-end model quality. +""" + +import argparse +import hashlib +import json +import statistics +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +import vllm._C_gguf # noqa: F401 + +import vllm._custom_ops # noqa: F401 +from vllm.model_executor.layers.quantization.gguf_expert_cache import ( + TurboMindPrefill, + _canonical, +) +from vllm.model_executor.layers.quantization.gguf_native import packed_tp_span +from vllm.transformers_utils.gguf_tensor_reader import ( + GGUFReader, + dequantize, + quant_size, +) + +EXPERTS, TOPK, HIDDEN, INTERMEDIATE = 512, 10, 2560, 160 + + +def artifact_identity(): + import vllm._C + + import vllm + + native = Path(vllm._C.__file__).resolve() + return { + "package": str(Path(vllm.__file__).resolve()), + "native": str(native), + "native_sha256": hashlib.sha256(native.read_bytes()).hexdigest(), + } + + +def measure(fn, iterations): + for _ in range(3): + fn() + torch.accelerator.synchronize() + samples = [] + for _ in range(5): + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + samples.append(begin.elapsed_time(end) / iterations) + return {"median_ms": statistics.median(samples), "samples_ms": samples} + + +def pointer_table(tensor, stride, ld): + table = torch.empty(EXPERTS, 2, dtype=torch.int64, device="cuda") + table[:, 0] = tensor.data_ptr() + torch.arange(EXPERTS, device="cuda") * stride + table[:, 1] = ld + return table.view(torch.uint8).flatten() + + +def relative_l2(actual, expected): + actual, expected = actual.float(), expected.float() + return ( + ((actual - expected).norm(dim=-1) / expected.norm(dim=-1).clamp_min(1e-6)) + .max() + .item() + ) + + +def read_layer(reader, layer): + tensors = { + part: next( + t for t in reader.tensors if t.name == f"blk.{layer}.ffn_{part}_exps.weight" + ) + for part in ("gate", "up", "down") + } + gtype, dtype = (int(tensors[p].tensor_type) for p in ("gate", "down")) + assert int(tensors["up"].tensor_type) == gtype + raw = { + p: np.ascontiguousarray(np.asarray(tensors[p].data)[:, :INTERMEDIATE]) + for p in ("gate", "up") + } + full = np.asarray(tensors["down"].data) + block, size = quant_size(dtype) + first, last, left, _ = packed_tp_span(full.shape[2] // size * block, dtype, 4, 0) + raw["down"] = np.ascontiguousarray(full[:, :, first:last]) + return raw, gtype, dtype, left + + +def load_layer(reader, layer): + raw, gtype, dtype, left = read_layer(reader, layer) + source = {p: torch.from_numpy(a).cuda() for p, a in raw.items()} + info = { + "gate": (gtype,), + "down": (dtype,), + "rows": INTERMEDIATE, + "hidden": HIDDEN, + "gu": _canonical(gtype, 2 * INTERMEDIATE, HIDDEN), + "dn": _canonical(dtype, HIDDEN, INTERMEDIATE), + } + banks, tables = {}, {} + ids = torch.arange(EXPERTS, dtype=torch.int32, device="cuda") + count = torch.full((1,), EXPERTS, dtype=torch.int32, device="cuda") + error = torch.zeros(1, dtype=torch.int32, device="cuda") + pointers = { + p: t.data_ptr() + + torch.arange(EXPERTS, device="cuda", dtype=torch.int64) * t[0].numel() + for p, t in source.items() + } + for key, part, n, k, skip in ( + ("gu", "gate", 2 * INTERMEDIATE, HIDDEN, 0), + ("dn", "down", HIDDEN, INTERMEDIATE, left), + ): + _, _, _, weight_bytes, stats_bytes = info[key] + weight = torch.empty(EXPERTS * weight_bytes, dtype=torch.uint8, device="cuda") + stats = torch.empty(EXPERTS * stats_bytes, dtype=torch.uint8, device="cuda") + banks[key] = (weight, stats) + tables[key] = ( + pointer_table(weight, weight_bytes, 32 * k), + pointer_table(stats, stats_bytes, n), + ) + torch.ops._C.gguf_prefill_transcode_sm70( + weight, + stats, + weight_bytes, + stats_bytes, + pointers[part], + pointers["up"] if key == "gu" else None, + ids, + count, + 1024, + gtype if key == "gu" else dtype, + n, + k, + skip, + raw[part].shape[2], + error, + ) + assert error.item() == 0, "transcode reported an invalid expert" + return info, tables, banks, raw, left + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--layers", type=int, nargs="+", default=[0, 1, 17, 47]) + parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347]) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--tune-tokens", type=int, default=8192) + parser.add_argument("--down-tokens", type=int, default=4096) + parser.add_argument("--routing", choices=["uniform", "skewed"], default="uniform") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--save-reference", type=Path) + parser.add_argument("--compare-reference", type=Path) + args = parser.parse_args() + assert args.iterations > 0 and min(args.tokens) > 0 and args.down_tokens > 0 + torch.backends.cuda.matmul.allow_tf32 = False + reader = GGUFReader(args.model) + references = ( + torch.load(args.compare_reference, map_location="cpu", weights_only=True) + if args.compare_reference + else {} + ) + saved, results = {}, [] + report = { + "artifact": artifact_identity(), + "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(), + "routing": args.routing, + "tune_tokens": args.tune_tokens, + "down_tokens": args.down_tokens, + "cases": results, + } + for layer in args.layers: + info, tables, banks, raw, left = load_layer(reader, layer) + # Reuse the production sorting, sub-batching, fused GEMM, and unroute. + runner = TurboMindPrefill.__new__(TurboMindPrefill) + runner.cache = SimpleNamespace(experts=EXPERTS, prefill_tokens=args.down_tokens) + runner.expert_ids = torch.arange(EXPERTS + 1, dtype=torch.int32, device="cuda") + runner.tuned = set() + probe = torch.zeros(args.tune_tokens, HIDDEN, dtype=torch.half, device="cuda") + runner._tune(layer, info, tables, probe) + del probe + dequantized = {} + + def ref_weight( + expert, part, dequantized=dequantized, info=info, raw=raw, left=left + ): + key = (expert, part) + if key not in dequantized: + ty = info["down" if part == "down" else "gate"][0] + array = dequantize(raw[part][expert], ty) + if part == "down": + array = array[:, left : left + INTERMEDIATE] + dequantized[key] = ( + torch.from_numpy(np.ascontiguousarray(array)).cuda().float() + ) + return dequantized[key] + + for tokens in args.tokens: + gen = torch.Generator(device="cuda").manual_seed(tokens) + x = (torch.randn(tokens, HIDDEN, device="cuda", generator=gen) * 0.5).half() + scores = torch.randn(tokens, EXPERTS, device="cuda", generator=gen) + if args.routing == "skewed": + scores[:, :32] += 2 + scores, ids = scores.topk(TOPK, dim=1) + weights = scores.softmax(dim=1) + route = runner._route(ids, 0, EXPERTS) + subroutes = ( + [ + ( + start, + min(tokens, start + args.down_tokens), + runner._subroute( + route, start, min(tokens, start + args.down_tokens) + ), + ) + for start in range(0, tokens, args.down_tokens) + ] + if tokens > args.down_tokens + else [(0, tokens, route)] + ) + out = torch.empty_like(x) + + def gu(runner=runner, info=info, tables=tables, x=x, route=route): + return runner._gate_up(info, tables, x, route, False) + + hidden = gu() + + def down( + h=hidden, + runner=runner, + info=info, + tables=tables, + subroutes=subroutes, + weights=weights, + out=out, + ): + for start, end, subroute in subroutes: + runner._down( + info, + tables, + h, + subroute, + weights[start:end], + out[start:end], + False, + ) + + def chain(gu=gu, down=down): + down(gu()) + + result = { + "layer": layer, + "gate_type": info["gate"][0], + "down_type": info["down"][0], + "tokens": tokens, + "gate_up": measure(gu, args.iterations), + "down_unroute": measure(down, args.iterations), + "chain": measure(chain, args.iterations), + } + down() + # FP32 reference at expert boundaries catches ragged-row/layout errors. + offsets = route[0].cpu().tolist() + positions = sorted( + { + p + for e in (0, 1, 31, 32, 127, 255, 510, 511) + if offsets[e] < offsets[e + 1] + for p in (offsets[e], offsets[e + 1] - 1) + } + ) + hidden_errors = [] + for position in positions: + token = int(route[2][position]) + expert = int(ids.flatten()[route[1][position]]) + expected = torch.nn.functional.silu( + ref_weight(expert, "gate") @ x[token].float() + ) + expected *= ref_weight(expert, "up") @ x[token].float() + hidden_errors.append(relative_l2(hidden[position], expected)) + chosen_tokens = torch.tensor( + sorted({0, 1, tokens // 2, tokens - 1}), device="cuda" + ).clamp_max(tokens - 1) + sample = { + "hidden": hidden[positions].cpu(), + "out": out[chosen_tokens].cpu(), + } + expected_rows = [] + for token in chosen_tokens.tolist(): + expected_row = torch.zeros(HIDDEN, device="cuda") + for slot in range(TOPK): + expert = int(ids[token, slot]) + expected_hidden = torch.nn.functional.silu( + ref_weight(expert, "gate") @ x[token].float() + ) * (ref_weight(expert, "up") @ x[token].float()) + expected_down = ( + (ref_weight(expert, "down") @ expected_hidden.half().float()) + .half() + .float() + ) + expected_row += weights[token, slot] * expected_down + expected_rows.append(expected_row.half()) + output_error = relative_l2(out[chosen_tokens], torch.stack(expected_rows)) + key = f"{layer}/{tokens}/{args.routing}" + result["hidden_fp32_max_rel_l2"] = max(hidden_errors) + result["output_fp32_max_rel_l2"] = output_error + result["finite"] = bool( + torch.isfinite(hidden).all() & torch.isfinite(out).all() + ) + passed = ( + result["finite"] and max(hidden_errors) <= 0.01 and output_error <= 0.01 + ) + if args.compare_reference: + errors = { + name: relative_l2(value, references[key][name]) + for name, value in sample.items() + } + result["control_max_rel_l2"] = errors + passed &= max(errors.values()) <= 0.003 + result["passed"] = passed + saved[key] = sample + results.append(result) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(result), flush=True) + del banks, tables, dequantized, hidden, out, runner + if args.save_reference: + torch.save(saved, args.save_reference) + if not all(case["passed"] for case in results): + raise SystemExit("numerical gate failed") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/mtp_trace_worker.py b/benchmarks/flashnext_prefill/mtp_trace_worker.py new file mode 100644 index 0000000000..d2748b5dcc --- /dev/null +++ b/benchmarks/flashnext_prefill/mtp_trace_worker.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diagnostic NVTX annotations installed after graph capture. + +Use only for the Nsight composition pass. Ordinary speed admission uses the +normal worker, without these wrappers or synchronized step profiling. +""" + +from functools import wraps + +import torch + +from vllm.v1.worker.gpu_worker import Worker + + +def annotate(owner, name, label): + original = getattr(owner, name) + + @wraps(original) + def observed(*args, **kwargs): + text = label(*args, **kwargs) if callable(label) else label + with torch.cuda.nvtx.range(text): + return original(*args, **kwargs) + + setattr(owner, name, observed) + + +class TraceWorker(Worker): + def profile(self, is_start=True, profile_prefix=None): + if is_start: + # The benchmark reaches this boundary with an idle request queue. + # Release unused eager-prefill allocations before CUPTI allocates + # its buffers; graph-owned tensors and pointers remain live. + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + return super().profile(is_start, profile_prefix) + + def compile_or_warm_up_model(self): + result = super().compile_or_warm_up_model() + runner = self.model_runner + prepared = {} + prepare = runner.prepare_inputs + + @wraps(prepare) + def prepare_inputs(*args, **kwargs): + with torch.cuda.nvtx.range("flashnext_batch.prepare"): + batch = prepare(*args, **kwargs) + prepared.update( + requests=batch.num_reqs, + tokens=batch.num_tokens, + drafts=batch.num_draft_tokens, + ) + return batch + + runner.prepare_inputs = prepare_inputs + annotate(runner, "sample", "flashnext_batch.target_sample") + annotate(runner, "postprocess_sampled", "flashnext_batch.target_state") + annotate( + runner.cudagraph_manager, + "run_fullgraph", + lambda desc: ( + f"flashnext_batch.target/M{desc.num_tokens}/R{desc.num_reqs}" + f"/actualR{prepared['requests']}/T{prepared['tokens']}" + f"/D{prepared['drafts']}" + ), + ) + draft = runner.speculator + annotate( + draft, + "propose", + lambda batch, *args, **kwargs: f"flashnext_batch.draft/R{batch.num_reqs}", + ) + for manager, name in ( + (draft.prefill_cudagraph_manager, "draft_first"), + (draft.decode_cudagraph_manager, "draft_rest"), + (getattr(draft, "multistep_cudagraph_manager", None), "draft_chain"), + ): + if manager is None: + continue + annotate( + manager, + "run_fullgraph", + lambda desc, name=name: f"flashnext_batch.{name}/M{desc.num_tokens}", + ) + return result diff --git a/benchmarks/flashnext_prefill/nsys_breakdown.py b/benchmarks/flashnext_prefill/nsys_breakdown.py index 7af351a9c2..7d0e3270e2 100644 --- a/benchmarks/flashnext_prefill/nsys_breakdown.py +++ b/benchmarks/flashnext_prefill/nsys_breakdown.py @@ -14,22 +14,27 @@ c = sqlite3.connect(sys.argv[1]) dev = int(sys.argv[2]) names = dict(c.execute("select id, value from StringIds").fetchall()) -api = dict( - c.execute("select correlationId, end from CUPTI_ACTIVITY_KIND_RUNTIME").fetchall() -) +# CUDA correlation ids are process-local. Nsight packs the process id above +# the low 24 thread-id bits; joining only correlationId mixes TP workers. +api = { + (tid >> 24, cid): end + for tid, cid, end in c.execute( + "select globalTid, correlationId, end from CUPTI_ACTIVITY_KIND_RUNTIME" + ) +} ks = c.execute( - "select start,end,streamId,correlationId,demangledName,gridX,gridY " + "select start,end,streamId,correlationId,demangledName,gridX,gridY,globalPid " "from CUPTI_ACTIVITY_KIND_KERNEL where deviceId=?", (dev,), ).fetchall() ms = c.execute( - "select start,end,streamId,correlationId,NULL,0,0 " + "select start,end,streamId,correlationId,NULL,0,0,globalPid " "from CUPTI_ACTIVITY_KIND_MEMCPY where deviceId=?", (dev,), ).fetchall() -st = collections.Counter(r[2] for r in ks) +st = collections.Counter((r[-1], r[2]) for r in ks) main = st.most_common(1)[0][0] -m = sorted(r for r in ks + ms if r[2] == main) +m = sorted(r for r in ks + ms if (r[-1], r[2]) == main) phase = "gu" @@ -128,11 +133,12 @@ def seg(n, gx, gy): wait = collections.defaultdict(float) host = 0 end = m[0][0] -for s, e, _, cid, nm, gx, gy in m: +for s, e, _, cid, nm, gx, gy, pid in m: name = names.get(nm) if nm is not None else None g = seg(name, gx, gy) if s > end: - if api.get(cid) is not None and api[cid] <= end: + submitted = api.get((pid >> 24, cid)) + if submitted is not None and submitted <= end: wait[g] += s - end else: host += s - end @@ -151,9 +157,9 @@ def seg(n, gx, gy): f"wait-before {wait[g] / 1e6:5.0f} n={n[g]}" ) side = collections.defaultdict(float) -for s, e, sid, cid, nm, gx, gy in ks: - if sid != main: - side[(sid, (names.get(nm) or "")[:40])] += e - s +for s, e, sid, cid, nm, gx, gy, pid in ks: + if (pid, sid) != main: + side[(pid >> 24, sid, (names.get(nm) or "")[:40])] += e - s print( "side streams:", [(k, round(v / 1e6)) for k, v in sorted(side.items(), key=lambda x: -x[1])[:5]], diff --git a/benchmarks/flashnext_prefill/quality.py b/benchmarks/flashnext_prefill/quality.py new file mode 100644 index 0000000000..52a237e907 --- /dev/null +++ b/benchmarks/flashnext_prefill/quality.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired retrieval/arithmetic and EOS checks at specified prompt-token budgets. + +Uses production sampling and thinking. Reset only an exclusively owned server. +Record the same prompt hash, seed and actual token length for each artifact. +""" + +import argparse +import hashlib +import json +import time +from collections import Counter +from pathlib import Path + +import regex as re +from measure_prefill import request + +FILLER = ( + "Computing records describe mechanical calculators, vacuum tubes, transistors, " + "integrated circuits, storage systems and communication networks. " + "These background records do not contain any vault or captain information.\n" +) +EXPECTED = {"vault": "48213", "captain": "Orlanda Vexley", "packages": 45} + + +def prompt(repeats): + return ( + "Read this archive and answer the question at the end.\n" + + FILLER * (repeats // 8) + + "The vault code is 48213.\n" + + FILLER * (repeats // 2) + + "The ship captain is Orlanda Vexley.\n" + + FILLER * (repeats - repeats // 8 - repeats // 2) + + "A warehouse has 17 boxes with 24 items each. After selling 93 items, " + "all remaining items are packed into groups of 7. How many full packages " + "are made? Return only one JSON object containing the vault code as a " + 'string, captain name, and number of packages: {"vault": ..., ' + '"captain": ..., "packages": ...}. Think briefly.\n' + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8002") + parser.add_argument("--budgets", type=int, nargs="+", default=[32768]) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if min(args.budgets) < 512: + parser.error("prompt budgets must be at least 512 tokens") + results = [] + for budget in args.budgets: + + def count(repeats): + messages = [{"role": "user", "content": prompt(repeats)}] + with request( + args.base_url, + "/tokenize", + { + "model": "flash-next", + "messages": messages, + "chat_template_kwargs": {"enable_thinking": True}, + }, + ) as response: + return json.load(response)["count"], messages + + fixed, _ = count(0) + ten, _ = count(10) + repeats = max(0, int((budget - fixed) * 10 / (ten - fixed))) + tokens, messages = count(repeats) + while tokens > budget: + repeats -= 1 + tokens, messages = count(repeats) + body = { + "model": "flash-next", + "messages": messages, + "max_tokens": 1024, + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "seed": 9683, + "chat_template_kwargs": {"enable_thinking": True}, + } + with request(args.base_url, "/reset_prefix_cache", {}) as response: + response.read() + start = time.perf_counter() + with request(args.base_url, "/v1/chat/completions", body) as response: + result = json.load(response) + choice = result["choices"][0] + message = choice["message"] + content = message.get("content") or "" + match = re.search(r"\{[^{}]*\}", content) + try: + answer = json.loads(match.group()) if match else None + except json.JSONDecodeError: + answer = None + text = " ".join( + str(message.get(k) or "") + for k in ("content", "reasoning", "reasoning_content") + ) + words = text.split() + ngrams = Counter(tuple(words[i : i + 10]) for i in range(len(words) - 9)) + max_repeat = max(ngrams.values(), default=0) + row = { + "budget_tokens": budget, + "tokenize_count": tokens, + "prompt_sha256": hashlib.sha256( + messages[0]["content"].encode() + ).hexdigest(), + "seconds": time.perf_counter() - start, + "usage": result["usage"], + "choice": choice, + "max_ten_word_repeat": max_repeat, + "passed": answer == EXPECTED + and choice["finish_reason"] == "stop" + and max_repeat <= 5, + } + results.append(row) + args.output.write_text(json.dumps(results, indent=2, ensure_ascii=False) + "\n") + print(json.dumps(row, ensure_ascii=False), flush=True) + raise SystemExit(0 if all(row["passed"] for row in results) else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py new file mode 100644 index 0000000000..b140fb8a9a --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""TP4 paired screens of the existing expert-cache byte-copy kernel.""" + +import argparse +import json +import os +import statistics +from functools import partial +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm import _custom_ops # noqa: F401 -- register mapped-view operator +from vllm.model_executor.layers.quantization.gguf_expert_cache import _apply_plan + +POLICIES = { + "reference": (16, 8, 1024, 8, False), + "o2s64b256w4": (2, 64, 256, 4, False), + "o4s32b256w4": (4, 32, 256, 4, False), + "o4s64b256w4": (4, 64, 256, 4, False), + "o8s16b512w4": (8, 16, 512, 4, False), + "o2s32b512w4": (2, 32, 512, 4, False), + "o2s32b256w4p3": (2, 32, 256, 4, True), + "o4s32b256w4p3": (4, 32, 256, 4, True), + "o4s64b256w4p3": (4, 64, 256, 4, True), +} + + +def launch(buffers, policy): + programs, split, block, warps, parallel_shards = POLICIES[policy] + plan, arena, mapped, device_offsets, host_offsets, words = buffers + grid = (programs, split, 3) if parallel_shards else (programs, split) + _apply_plan[grid]( + plan, + arena, + mapped, + device_offsets, + host_offsets, + words, + 0, + OPS_PROGRAMS=programs, + SPLIT=split, + BLOCK=block, + PARALLEL_SHARDS=parallel_shards, + num_warps=warps, + ) + + +def reset_data(arena, host, original_arena, original_host): + arena.copy_(original_arena) + host.copy_(original_host) + + +def summarize(values): + ordered = sorted(values) + return { + "mean": statistics.mean(values), + "p50": statistics.median(values), + "p90": ordered[round(0.9 * (len(ordered) - 1))], + "values": values, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--samples", type=int, default=12) + parser.add_argument("--repeats", type=int, default=16) + parser.add_argument("--boundary-only", action="store_true") + parser.add_argument("--policies", nargs="+", choices=POLICIES) + args = parser.parse_args() + policies = ( + {name: POLICIES[name] for name in args.policies} if args.policies else POLICIES + ) + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + device = torch.device("cuda", rank) + # Flash-Next GSQ-RCO IQ3_S TP4 words per expert. Q2_0's down projection + # retains its 192-column physical padding for the 160-column local input. + layouts = { + 0: dict(zip(("w1", "w3", "w2"), (19600, 19600, 28800))), + 1: dict(zip(("w1", "w3", "w2"), (16400, 16400, 17280))), + 17: dict(zip(("w1", "w3", "w2"), (22000, 22000, 28800))), + 47: dict(zip(("w1", "w3", "w2"), (27200, 27200, 28800))), + } + if args.shapes: + layouts = {} + for entry in json.loads(args.shapes.read_text()): + layer = int(entry["name"].split(".")[1]) + shard = ( + "w2" + if "down" in entry["name"] + else "w1" + if "gate" in entry["name"] + else "w3" + ) + assert entry["rank_bytes"] % 8 == 0 + layouts.setdefault(layer, {})[shard] = entry["rank_bytes"] // 8 + if args.boundary_only: + layouts = {-1: {"w1": 255, "w3": 256, "w2": 257}} + output = { + "world_size": 4, + "policies": policies, + "samples": args.samples, + "repeats": args.repeats, + "cases": [], + "correctness": [], + } + + for layer, shards in sorted(layouts.items()): + sizes = [shards[k] for k in ("w1", "w3", "w2")] + # Production slot counts; tail guards separate projection banks. + slots, host_slots = 410, 136 + offsets, host_offsets, total, host_total = [], [], 0, 0 + for n in sizes: + total = ((total + 31) // 32) * 32 + offsets.append(total) + host_offsets.append(host_total) + total += slots * n + 32 + host_total += host_slots * n + arena = torch.arange(total, dtype=torch.int64, device=device) + host = torch.empty(host_total, dtype=torch.int64, pin_memory=True) + host.copy_(torch.arange(host_total, dtype=torch.int64).mul_(-3).sub_(7)) + original_arena, original_host = arena.clone(), host.clone() + mapped = torch.ops._C.get_cuda_view_from_cpu_tensor(host) + assert mapped.device == device + device_offsets = torch.tensor(offsets, dtype=torch.int64, device=device) + cpu_offsets = torch.tensor(host_offsets, dtype=torch.int64, device=device) + words = torch.tensor(sizes, dtype=torch.int64, device=device) + plan = torch.zeros(8 + 5 * 256, dtype=torch.int32, device=device) + + buffers = (plan, arena, mapped, device_offsets, cpu_offsets, words) + invoke = partial(launch, buffers) + reset = partial(reset_data, arena, host, original_arena, original_host) + + # Separate loads/writebacks model the spare-slot policy; a swap is + # also valid when an incoming expert gives its host slot to a victim. + cases = ( + ("empty", 0), + ("paired", 1), + ("paired", 2), + ("paired", 4), + ("paired", 8), + ("paired", 20), + ("swap", 1), + ("swap", 4), + ) + if args.boundary_only: + cases = tuple(("swap", n) for n in (0, 1, 2, 3, 4, 5, 15, 16, 17, 128)) + for mode, count in cases: + cpu_plan = torch.zeros_like(plan, device="cpu") + ops = count * 2 if mode == "paired" else count + cpu_plan[0] = ops + for i in range(count): + entries = ( + [(i, i, i)] + if mode == "swap" + else [(i, i, -1), (slots - count + i, -1, host_slots - count + i)] + ) + for j, (slot, load, back) in enumerate(entries): + op = i * len(entries) + j + cpu_plan[8 + op * 5 : 11 + op * 5] = torch.tensor( + [slot, load, back] + ) + plan.copy_(cpu_plan) + reset() + invoke("reference") + torch.accelerator.synchronize() + oracle_arena, oracle_host = arena.clone(), host.clone() + graphs = {} + for name in policies: + reset() + invoke(name) + torch.accelerator.synchronize() + assert torch.equal(arena, oracle_arena), ( + rank, + layer, + mode, + count, + name, + "GPU", + ) + assert torch.equal(host, oracle_host), ( + rank, + layer, + mode, + count, + name, + "host", + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(args.repeats): + invoke(name) + graphs[name] = graph + reset() + graph.replay() + torch.accelerator.synchronize() + # Even swap counts restore the original data after a graph. + expected_arena = ( + original_arena + if mode == "swap" and args.repeats % 2 == 0 + else oracle_arena + ) + expected_host = ( + original_host + if mode == "swap" and args.repeats % 2 == 0 + else oracle_host + ) + assert torch.equal(arena, expected_arena), (rank, name, "graph GPU") + assert torch.equal(host, expected_host), (rank, name, "graph host") + reset() + arena.add_(19) + host.add_(23) + for _ in range(args.repeats): + invoke("reference") + torch.accelerator.synchronize() + changed_arena, changed_host = arena.clone(), host.clone() + reset() + arena.add_(19) + host.add_(23) + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(arena, changed_arena), (rank, name, "changed GPU") + assert torch.equal(host, changed_host), (rank, name, "changed host") + del changed_arena, changed_host + + timings = {name: [] for name in policies} + for sample in range(args.samples): + order = list(policies) + if sample % 2: + order.reverse() + for name in order: + reset() + dist.barrier() + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + graphs[name].replay() + end.record() + end.synchronize() + elapsed = torch.tensor( + start.elapsed_time(end) * 1000 / args.repeats, + dtype=torch.float64, + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX) + timings[name].append(elapsed.item()) + case = { + "layer": layer, + "words": sizes, + "mode": mode, + "count": count, + "plan_ops": ops, + "critical_rank_us": { + name: summarize(values) for name, values in timings.items() + }, + } + output["cases"].append(case) + output["correctness"].append( + { + "layer": layer, + "mode": mode, + "count": count, + "eager_and_graph_bit_exact": True, + } + ) + if rank == 0: + print(json.dumps(case), flush=True) + args.output.write_text(json.dumps(output, indent=2)) + del graphs, oracle_arena, oracle_host + del invoke, reset, buffers, arena, host, mapped, original_arena, original_host + torch.accelerator.empty_cache() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/csrc/gguf_expert_cache.cpp b/csrc/gguf_expert_cache.cpp index 012e11249e..a3310326d6 100644 --- a/csrc/gguf_expert_cache.cpp +++ b/csrc/gguf_expert_cache.cpp @@ -257,6 +257,43 @@ class ExpertCacheLane { return extent; } + // Stage only host slots used by the most recent residency snapshot. Keep + // their offsets unchanged: the GPU source map still addresses those slots. + // The caller orders the copy after decode write-backs and the preceding + // cold transcode, exactly as for a full-arena copy. + void stage_cold(int64_t layer, int64_t shard, int64_t destination, + int64_t stream_ptr) { + TORCH_CHECK(layer >= 0 && layer < layers_ && shard >= 0 && shard < kShards, + "expert cache: bad cold stage bank"); + auto stream = reinterpret_cast(stream_ptr); + std::lock_guard guard(mutex_); + const int32_t* map = residency_maps_ + 2 * layer * experts_; + std::vector occupied(host_slots_, 0); + for (int64_t e = 0; e < experts_; ++e) { + if (map[e] >= 0) continue; + const int32_t h = map[experts_ + e]; + TORCH_CHECK(h >= 0 && h < host_slots_, + "expert cache: non-resident expert has no host slot"); + occupied[h] = 1; + } + check(cudaSetDevice(static_cast(device_)), "set device"); + const int64_t stride = bytes(layer, shard); + for (int64_t begin = 0; begin < host_slots_;) { + if (!occupied[begin]) { + ++begin; + continue; + } + int64_t end = begin + 1; + while (end < host_slots_ && occupied[end]) ++end; + check(cudaMemcpyAsync( + reinterpret_cast(destination) + begin * stride, + host(layer, shard, begin), (end - begin) * stride, + cudaMemcpyHostToDevice, stream), + "stage cold experts"); + begin = end; + } + } + private: int64_t bytes(int64_t layer, int shard) const { return slot_bytes_[layer * kShards + shard]; @@ -533,6 +570,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("prefill", &ExpertCacheLane::prefill, pybind11::call_guard()) .def("residency", &ExpertCacheLane::residency, + pybind11::call_guard()) + .def("stage_cold", &ExpertCacheLane::stage_cold, pybind11::call_guard()); m.def("host_register", &host_register); m.def("host_unregister", &host_unregister); diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index d25941f769..2b04d8c1c6 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -366,6 +366,12 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X, int64_t top_k, int64_t type, int64_t row, int64_t tokens); +torch::stable::Tensor ggml_moe_a8_vec_rows(torch::stable::Tensor X, + torch::stable::Tensor W, + torch::stable::Tensor topk_ids, + int64_t top_k, int64_t type, + int64_t row, int64_t tokens); + int64_t ggml_moe_get_block_size(int64_t type); void paged_attention_v1( diff --git a/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu b/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu index 0fdfcafab8..23225b6b37 100644 --- a/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu +++ b/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu @@ -393,11 +393,11 @@ torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X, // input return Y; } -torch::stable::Tensor ggml_moe_a8_vec( +static torch::stable::Tensor ggml_moe_a8_vec_impl( torch::stable::Tensor X, // input torch::stable::Tensor W, // expert weights torch::stable::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row, - int64_t tokens) { + int64_t tokens, int rows_per_cta) { int col = X.sizes()[1]; const int padded = (col + 512 - 1) / 512 * 512; const torch::stable::accelerator::DeviceGuard device_guard( @@ -442,7 +442,7 @@ torch::stable::Tensor ggml_moe_a8_vec( moe_vec_q8_0_q8_1_cuda( (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); + col, row, quant_X.stride(0), stream, rows_per_cta); break; case 10: moe_vec_q2_K_q8_1_cuda( @@ -533,6 +533,43 @@ torch::stable::Tensor ggml_moe_a8_vec( return Y; } +torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X, + torch::stable::Tensor W, + torch::stable::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, + int64_t tokens) { + return ggml_moe_a8_vec_impl(X, W, topk_ids, top_k, type, row, tokens, + GGML_CUDA_MMV_Y); +} + +torch::stable::Tensor ggml_moe_a8_vec_rows(torch::stable::Tensor X, + torch::stable::Tensor W, + torch::stable::Tensor topk_ids, + int64_t top_k, int64_t type, + int64_t row, int64_t tokens) { + // One independent warp per output row; quantization, dp4a and warp sums + // are identical to the retained one-row CTA. Only CTA scheduling changes. + STD_TORCH_CHECK( + type == 8 && X.scalar_type() == torch::headeronly::ScalarType::Half, + "Row-batched GGUF MoE requires FP16 inputs and Q8_0 weights"); + STD_TORCH_CHECK( + X.dim() == 2 && W.dim() == 3 && topk_ids.dim() == 2 && + X.is_contiguous() && W.is_contiguous() && topk_ids.is_contiguous() && + W.scalar_type() == torch::headeronly::ScalarType::Byte && + X.size(1) % 32 == 0 && W.size(1) == row && + W.size(2) == X.size(1) / 32 * 34 && tokens == X.size(0) && + top_k > 0 && topk_ids.numel() == tokens * top_k && + topk_ids.scalar_type() == torch::headeronly::ScalarType::Int, + "Invalid row-batched GGUF MoE layout"); + STD_TORCH_CHECK(X.device() == W.device() && X.device() == topk_ids.device(), + "Row-batched GGUF MoE requires same-device tensors"); + if (tokens == 0) { + return torch::stable::empty({0, row}, X.scalar_type(), std::nullopt, + W.device()); + } + return ggml_moe_a8_vec_impl(X, W, topk_ids, top_k, type, row, tokens, 8); +} + int64_t ggml_moe_get_block_size(int64_t type) { switch (type) { case 2: diff --git a/csrc/libtorch_stable/quantization/gguf/moe.cuh b/csrc/libtorch_stable/quantization/gguf/moe.cuh index a2f9f46c8f..eed248f957 100644 --- a/csrc/libtorch_stable/quantization/gguf/moe.cuh +++ b/csrc/libtorch_stable/quantization/gguf/moe.cuh @@ -30,7 +30,9 @@ static __device__ __forceinline__ void moe_q( } const int exp_idx = expert_ids[blockIdx.y]; - if (exp_idx > 255 || exp_idx < 0) return; + // Negative IDs mark padding. Valid expert IDs may exceed 255 (e.g. 512 + // experts); skipping those leaves their output rows uninitialized. + if (exp_idx < 0) return; if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return; const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride); diff --git a/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh b/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh index 60f65a1bfd..108a940718 100644 --- a/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh +++ b/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh @@ -117,11 +117,11 @@ static void moe_vec_q8_0_q8_1_cuda(const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, const int top_k, const int tokens, const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const int token_stride, cudaStream_t stream, + const int rows_per_cta = GGML_CUDA_MMV_Y) { + const int block_num_y = (nrows + rows_per_cta - 1) / rows_per_cta; const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + const dim3 block_dims(WARP_SIZE, rows_per_cta, 1); moe_vec_q<<>>( vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 66a6ef88a4..c2a22e27f4 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -466,6 +466,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor topk_ids, int top_k, " "int type, SymInt row, SymInt tokens) -> Tensor"); + ops.def( + "ggml_moe_a8_vec_rows(Tensor X, Tensor W, Tensor topk_ids, int top_k, " + "int type, SymInt row, SymInt tokens) -> Tensor"); + ops.def("ggml_moe_get_block_size(int type) -> int"); // Mamba selective scan kernel @@ -623,6 +627,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8)); ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8)); ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec)); + ops.impl("ggml_moe_a8_vec_rows", TORCH_BOX(&ggml_moe_a8_vec_rows)); ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1)); diff --git a/csrc/ops.h b/csrc/ops.h index fd1dd131f2..c7db0d49fc 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1045,6 +1045,7 @@ void gguf_dp4a_gate_up_sm70_out(torch::Tensor out, torch::Tensor activation, torch::Tensor up, int64_t source_type, bool activated, int64_t lanes_per_row, bool bank_aware); + void gguf_dp4a_lut4_gate_up_sm70_out( torch::Tensor out, torch::Tensor activation, torch::Tensor ids, torch::Tensor gate, torch::Tensor gate_stats, torch::Tensor up, @@ -1056,6 +1057,12 @@ void gguf_dense_segments_sm70_out( std::vector out, std::vector fmt, std::vector n, int64_t k, int64_t split, int64_t warps, torch::Tensor ws, torch::Tensor cnt, std::optional sgate); +void gguf_dense_segments_batch_sm70_out( + torch::Tensor x, std::vector codes, + std::vector high, std::vector scale, + std::vector out, std::vector fmt, + std::vector n, int64_t k, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, std::optional sgate); void gguf_shared_gate_up_sm70_out(torch::Tensor x, std::vector gate, std::vector up, @@ -1186,3 +1193,48 @@ void gguf_dmvq_sm70_out(torch::Tensor output, torch::Tensor input, torch::Tensor counters, torch::Tensor table, int64_t type, int64_t kw, int64_t split); void gguf_dmvq_book_sm70_out(torch::Tensor table, int64_t type); + +void gguf_dp4a_gate_up_grouped_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode); + +void gguf_dp4a_gate_up_grouped_batch_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode); + +void gguf_dp4a_raw_down_grouped_sm70_out( + torch::Tensor out, torch::Tensor projected, torch::Tensor input, + torch::Tensor ids, torch::Tensor order, torch::Tensor probabilities, + torch::Tensor weights, int64_t source_type, int64_t left); + +void sm70_hcx_batch_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule); + +void sm70_hcx_batch_packed_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule); diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu index f4a632a5ff..41be11393b 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu @@ -18,6 +18,19 @@ using D = cache_policy::Default; namespace { +// Reuse the existing decoder and indexed grouped layout, with the padded +// activation supply already used by the quantized batch family. +template +struct PaddedActivationConfig; + +template +struct PaddedActivationConfig> { + using type = Sm70_s884, TransformA, U, B, + TransformB, V, order_C, Tc, raster_order, group_axis>; +}; + // Keep shape-specific small-N tactics out of prefill and unrelated CUDA graph // shapes. template @@ -124,6 +137,31 @@ void Registry::sm70_884_4() { 1, Group, 64, 64>; Add(std::make_unique< GgufPrefillNarrowKernelImpl>()); + // Ragged expert batches often have fewer than 128 rows. Halving M + // reduces both masked MMA work and the accumulator register demand; + // retain the existing quantization-group-sized K iteration. + using Prefill64x64 = + typename G::template Type<64, 64, Group, 2, 1, 1, D, D, 2, true, + 1, Group, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + using P = typename PaddedActivationConfig::type; + // K32/K64 unpadded activation rows repeat shared-memory banks across + // the four distinct rows read by each eight-lane MMA load group. + // Eight half values of row padding spread those accesses apart. + // A full-width epilogue is required: N64 subtiles of an N128 CTA + // are smaller than the epilogue thread map's horizontal span. + // K64 also distinguishes these tactics in serialized tuning caches. + using Padded32x64K64 = + typename P::template Type<32, 64, 64, 2, 2, 1, D, D, 2, true, + 1, Group, 32, 64>; + using Padded64x64K64 = + typename P::template Type<64, 64, 64, 2, 2, 1, D, D, 2, true, + 1, Group, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); } Add>(); @@ -159,6 +197,22 @@ void Registry::sm70_884_4() { GroupSize, 64, 64>; Add(std::make_unique< GgufPrefillNarrowKernelImpl>()); + using Prefill64x64 = + typename G::template Type<64, 64, 32, 2, 1, 1, D, D, 2, true, 1, + GroupSize, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + using P = typename PaddedActivationConfig::type; + using Padded32x64K64 = + typename P::template Type<32, 64, 64, 2, 2, 1, D, D, 2, true, + 1, GroupSize, 32, 64>; + using Padded64x64K64 = + typename P::template Type<64, 64, 64, 2, 2, 1, D, D, 2, true, + 1, GroupSize, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); } Add>(); diff --git a/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh index e8c9cf4059..d048bd489d 100644 --- a/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh @@ -368,18 +368,39 @@ __global__ void __launch_bounds__(32 * W) template void launch(const Segs& segs, const half* x, int ldx, int M, int K, int S, - int G, int split, float* ws, int* cnt, int tiles, cudaStream_t st) { - // Small-M verification needs three token tiles, but each uses the same - // weights. Keep the original path for other shapes and split-K schedules. + int G, int split, float* ws, int* cnt, int tiles, cudaStream_t st, + bool batch) { + // Reuse one weight decode across verification token tiles. Unsupported + // shapes and split-K schedules keep the existing row-tiled path. int total_n = 0; for (int i = 0; i < segs.nseg; ++i) total_n += segs.s[i].n; if constexpr (W == 4 || W == 8) { - if (M == 20 && split == 1 && - ((K == 2560 && (total_n == 4096 || total_n == 3584)) || - (K == 1536 && total_n == 2560))) { - dense_weight_major - <<>>( - segs, x, ldx, M, K, S, G, split, ws, cnt); + const bool legacy = + M == 20 && ((K == 2560 && (total_n == 4096 || total_n == 3584)) || + (K == 1536 && total_n == 2560)); + const bool batched = + batch && M > 8 && M <= 40 && + ((K == 2560 && total_n > 320) || (K == 1536 && total_n == 2560)); + if (split == 1 && (legacy || batched)) { +#define LAUNCH_WEIGHT_MAJOR(TM) \ + dense_weight_major \ + <<>>( \ + segs, x, ldx, M, K, S, G, split, ws, cnt) + switch ((M + 7) / 8) { + case 2: + LAUNCH_WEIGHT_MAJOR(2); + break; + case 3: + LAUNCH_WEIGHT_MAJOR(3); + break; + case 4: + LAUNCH_WEIGHT_MAJOR(4); + break; + case 5: + LAUNCH_WEIGHT_MAJOR(5); + break; + } +#undef LAUNCH_WEIGHT_MAJOR C10_CUDA_KERNEL_LAUNCH_CHECK(); return; } @@ -410,15 +431,20 @@ void launch(const Segs& segs, const half* x, int ldx, int M, int K, int S, } // namespace // codes/high/scale: per segment planes; out: per segment [M, n] views. -void gguf_dense_segments_sm70_out( - torch::Tensor x, std::vector codes, - std::vector high, std::vector scale, - std::vector out, std::vector fmt, - std::vector n, int64_t K, int64_t split, int64_t warps, - torch::Tensor ws, torch::Tensor cnt, std::optional sgate) { - TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kHalf && x.dim() == 2 && - x.stride(1) == 1 && x.size(0) > 0 && x.size(0) <= 32, - "input must be CUDA FP16 [1..32, K] with contiguous rows"); +template +void gguf_dense_segments_impl(torch::Tensor x, std::vector codes, + std::vector high, + std::vector scale, + std::vector out, + std::vector fmt, std::vector n, + int64_t K, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, + std::optional sgate) { + TORCH_CHECK( + x.is_cuda() && x.scalar_type() == at::kHalf && x.dim() == 2 && + x.stride(1) == 1 && x.size(0) > 0 && x.size(0) <= (Batch ? 40 : 32), + Batch ? "input must be CUDA FP16 [1..40, K] with contiguous rows" + : "input must be CUDA FP16 [1..32, K] with contiguous rows"); const c10::cuda::CUDAGuard guard(x.device()); const auto* properties = at::cuda::getCurrentDeviceProperties(); TORCH_CHECK(properties->major == 7 && properties->minor == 0, @@ -492,7 +518,7 @@ void gguf_dense_segments_sm70_out( cnt.numel() >= tiles * ((x.size(0) + 7) / 8), "split workspace or counter storage is too small"); const int M = static_cast(x.size(0)); - TORCH_CHECK(M <= 32); + TORCH_CHECK(M <= 40); const int S = static_cast((K + 31) / 32); const int G = (S + 3) / 4; auto st = at::cuda::getCurrentCUDAStream(); @@ -502,31 +528,51 @@ void gguf_dense_segments_sm70_out( int* cp = cnt.data_ptr(); switch (warps) { case 1: - launch<1>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<1>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 2: - launch<2>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<2>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 4: - launch<4>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<4>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 16: - launch<16>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<16>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 5: - launch<5>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<5>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 6: - launch<6>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<6>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 10: - launch<10>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<10>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 12: - launch<12>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<12>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; default: - launch<8>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<8>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; } } + +void gguf_dense_segments_sm70_out( + torch::Tensor x, std::vector codes, + std::vector high, std::vector scale, + std::vector out, std::vector fmt, + std::vector n, int64_t K, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, std::optional sgate) { + gguf_dense_segments_impl(x, codes, high, scale, out, fmt, n, K, split, + warps, ws, cnt, sgate); +} + +void gguf_dense_segments_batch_sm70_out( + torch::Tensor x, std::vector codes, + std::vector high, std::vector scale, + std::vector out, std::vector fmt, + std::vector n, int64_t K, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, std::optional sgate) { + gguf_dense_segments_impl(x, codes, high, scale, out, fmt, n, K, split, + warps, ws, cnt, sgate); +} diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh index ce28e96624..b850830c1a 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh @@ -55,6 +55,93 @@ struct LatticeDot { return book[index]; } + struct Decoded { + int words[8]; + float scale; + int scale0, scale1; + }; + + __device__ static Decoded decode(const uint8_t* row, int group, + const uint32_t* book, + const uint32_t* masks) { + Decoded operand{}; + const uint8_t* b = row + (group / 8) * kBlockBytes; + const int sub = group % 8; + const uint32_t low = load_u32_2(b + 2 + sub * (Type == 22 ? 4 : 8)); + uint32_t high = 0, signs = 0; + if constexpr (Type != 22) high = load_u32_2(b + 6 + sub * 8); + if constexpr (Type == 18) + signs = load_u32_2(b + 66 + sub * 4); + else + signs = load_u32_2(b + (Type == 21 ? 74 : 34) + sub * 4); + const int qh = Type == 18 ? 0 : b[66 + sub]; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + int first, second, sign; + if constexpr (Type == 22) { + first = + (((low >> (octet * 8)) & 255) | (((qh >> (2 * octet)) & 3) << 8)) * + 2; + second = first + 1; + } else { + const uint32_t codes = octet < 2 ? low : high; + const int shift = (octet % 2) * 16; + first = ((codes >> shift) & 255); + second = ((codes >> (shift + 8)) & 255); + if constexpr (Type == 21) { + first |= ((qh >> (2 * octet)) & 1) << 8; + second |= ((qh >> (2 * octet + 1)) & 1) << 8; + } + } + if constexpr (Type == 18) { + sign = (signs >> (7 * octet)) & 127; + sign |= (__popc(sign) & 1) << 7; + } else { + sign = (signs >> (8 * octet)) & 255; + } + const uint32_t s0 = BankAware ? sign_mask(sign & 15) : masks[sign & 15]; + const uint32_t s1 = BankAware ? sign_mask(sign >> 4) : masks[sign >> 4]; + operand.words[2 * octet] = __vsub4(word(book, first) ^ s0, s0); + operand.words[2 * octet + 1] = __vsub4(word(book, second) ^ s1, s1); + } + operand.scale = __half2float(*reinterpret_cast(b)); + if constexpr (Type == 21) { + operand.scale0 = 1 + 2 * ((b[106 + sub / 2] >> (4 * (sub % 2))) & 15); + } else if constexpr (Type == 18) { + operand.scale0 = 1 + 2 * (signs >> 28); + } else { + operand.scale0 = 1 + 2 * (b[74 + sub] & 15); + operand.scale1 = 1 + 2 * (b[74 + sub] >> 4); + } + return operand; + } + + __device__ static float dot_decoded(const Decoded& operand, const Q8_1& x) { + const int* activation = reinterpret_cast(x.qs); + int sum0 = 0, sum1 = 0; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + if constexpr (Type == 22) { + if (octet >= 2) { + sum1 = __dp4a(operand.words[2 * octet], activation[2 * octet], sum1); + sum1 = __dp4a(operand.words[2 * octet + 1], activation[2 * octet + 1], + sum1); + continue; + } + } + sum0 = __dp4a(operand.words[2 * octet], activation[2 * octet], sum0); + sum0 = + __dp4a(operand.words[2 * octet + 1], activation[2 * octet + 1], sum0); + } + const float d = operand.scale * __low2float(x.ds); + if constexpr (Type == 21) + return d * float(sum0 * operand.scale0); + else if constexpr (Type == 18) + return d * (float(sum0) * float(operand.scale0) * .25f); + else + return d * float(sum0 * operand.scale0 + sum1 * operand.scale1) * .125f; + } + __device__ static float dot(const uint8_t* row, int group, const Q8_1& x, const uint32_t* book, const uint32_t* masks) { const uint8_t* b = row + (group / 8) * kBlockBytes; @@ -133,6 +220,61 @@ using IQ3SDot = LatticeDot<21>; // canonical reader, without expanding a second expert bank at load time. template struct OriginalIntegerDot { + struct Decoded { + int values[8]; + float scale; + }; + __device__ static Decoded decode(const uint8_t* row, int group) { + const uint8_t* b; + const uint8_t* codes; + float scale; + if constexpr (Type == 23) { + b = row + (group / 8) * 136; + const int sub = group % 8; + const uint16_t high = *reinterpret_cast(b + 2); + const int subscale = ((b[4 + sub / 2] >> (4 * (sub % 2))) & 15) | + (((high >> (2 * sub)) & 3) << 4); + // Match the canonical LUT reader's expanded FP16 coefficients. + scale = __half2float(__float2half_rn( + __half2float(*reinterpret_cast(b)) * (subscale - 32))); + codes = b + 8 + sub * 16; + } else { + b = row + (Type == 20 ? group : group / 2) * 18; + scale = __half2float(*reinterpret_cast(b)); + codes = b + 2; + } + Decoded decoded; + decoded.scale = scale; +#pragma unroll + for (int packet = 0; packet < 8; ++packet) { + uint32_t values = 0; + if constexpr (Type == 20 || Type == 23) { + const uint32_t bytes = + load_u32_2(codes + (packet % 4) * 4) >> ((packet / 4) * 4); + const uint32_t nibbles = (bytes & 15) | ((bytes >> 4) & 0xf0) | + ((bytes >> 8) & 0xf00) | + ((bytes >> 12) & 0xf000); + using Lut = turbomind::gemm::Transform_HMMA_SM70_Lut4<0>; + values = Lut::iq_values(nibbles) ^ 0x80808080U; + } else { + const uint32_t byte = b[2 + (group % 2) * 8 + packet]; + const uint32_t codes = (byte & 3) | ((byte << 6) & 0x300) | + ((byte << 12) & 0x30000) | + ((byte << 18) & 0x3000000); + values = __vsub4(codes, 0x01010101U); + } + decoded.values[packet] = static_cast(values); + } + return decoded; + } + __device__ static float dot_decoded(const Decoded& decoded, const Q8_1& x) { + const int* activation = reinterpret_cast(x.qs); + int sum = 0; +#pragma unroll + for (int packet = 0; packet < 8; ++packet) + sum = __dp4a(decoded.values[packet], activation[packet], sum); + return float(sum) * (decoded.scale * __low2float(x.ds)); + } static_assert(Type == 20 || Type == 23 || Type == 42); __device__ static float dot(const uint8_t* row, int group, const Q8_1& x) { const uint8_t* b; diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu index faa16955c5..e1e9df682a 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu +++ b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu @@ -23,12 +23,17 @@ __global__ void quantize_q8(Q8_1* out, const half* input, int k) { template + bool ScalarLut = false, bool BankAware = false, int FixedGroups = 0, + bool SparseFallback = false> __global__ void gate_up(void* output, const Q8_1* activation, const Index* ids, const uint8_t* gate, const uint8_t* up, int n, int k, int stride, int top_k, const StridedPtr* gate_stats = nullptr, - const StridedPtr* up_stats = nullptr) { + const StridedPtr* up_stats = nullptr, + const int* grouping_mode = nullptr) { + if constexpr (SparseFallback) { + if (grouping_mode[0] != 0) return; + } using Dot = std::conditional_t< Canonical, vllm::sm70_gguf::CanonicalIntegerDot, std::conditional_t, @@ -712,3 +717,7 @@ void gguf_dp4a_scalar_lut_gate_up_sm70_out( dispatch_scalar_lut_gate_up<22>(out, activation, ids, gate, up, activated); C10_CUDA_KERNEL_LAUNCH_CHECK(); } + +#include "gguf_grouped_gate_sm70.cuh" + +#include "gguf_grouped_down_sm70.cuh" diff --git a/csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh new file mode 100644 index 0000000000..43c0425867 --- /dev/null +++ b/csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +namespace { +template +__global__ void raw_down_grouped_routes(half* projected, const Q8_1* input, + const int* ids, const int* order, + const uint8_t* weights, int routes, + int n, int groups, int stride, + int left_groups, int experts) { + __shared__ int active, count, expert, original[4]; + extern __shared__ Q8_1 shared_x[]; + const int route = blockIdx.y; + if (threadIdx.x == 0) { + expert = ids[route]; + int first = route; + while (first > 0 && ids[first - 1] == expert) --first; + active = ((route - first) % 4 == 0) && expert >= 0 && expert < experts; + count = 1; + while (count < 4 && route + count < routes && ids[route + count] == expert) + ++count; + for (int q = 0; q < count; ++q) original[q] = order[route + q]; + } + __syncthreads(); + if (!active) return; + for (int i = threadIdx.x; i < count * groups * 9; i += blockDim.x) { + const int src = original[i / (groups * 9)]; + reinterpret_cast(shared_x)[i] = + reinterpret_cast( + input)[int64_t(src) * groups * 9 + i % (groups * 9)]; + } + __syncthreads(); + const int col = blockIdx.x * 128 + threadIdx.x; + if (col >= n) return; + const uint8_t* row = weights + (int64_t(expert) * n + col) * stride; + float acc[4] = {}; + for (int group = 0; group < groups; ++group) { + const auto decoded = vllm::sm70_gguf::OriginalIntegerDot::decode( + row, group + left_groups); +#pragma unroll + for (int q = 0; q < 4; ++q) + if (q < count) + acc[q] += vllm::sm70_gguf::OriginalIntegerDot::dot_decoded( + decoded, shared_x[q * groups + group]); + } +#pragma unroll + for (int q = 0; q < 4; ++q) + if (q < count) + projected[int64_t(original[q]) * n + col] = __float2half_rn(acc[q]); +} + +__global__ void raw_down_grouped_unroute(half* output, const half* projected, + const float* weights, int n, + int top_k) { + const int token = blockIdx.y; + const int col = blockIdx.x * 128 + threadIdx.x; + if (col >= n) return; + // Retain the four warp-owned subtotals and their original addition order. + float sum = 0.f; +#pragma unroll + for (int owner = 0; owner < 4; ++owner) { + float partial = 0.f; + for (int route = owner; route < top_k; route += 4) { + const int slot = token * top_k + route; + partial += + __half2float(projected[int64_t(slot) * n + col]) * weights[slot]; + } + sum += partial; + } + output[int64_t(token) * n + col] = __float2half_rn(sum); +} + +template +void launch_grouped_raw_down(torch::Tensor out, torch::Tensor projected, + torch::Tensor input, torch::Tensor ids, + torch::Tensor order, torch::Tensor probabilities, + torch::Tensor weights, int left) { + const int n = out.size(1), top_k = ids.size(1), groups = input.size(2); + const auto stream = at::cuda::getCurrentCUDAStream(); + raw_down_grouped_routes<<>>( + reinterpret_cast(projected.data_ptr()), + reinterpret_cast(input.data_ptr()), ids.data_ptr(), + order.data_ptr(), weights.data_ptr(), ids.numel(), n, + groups, weights.size(2), left / 32, weights.size(0)); + raw_down_grouped_unroute<<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(projected.data_ptr()), + probabilities.data_ptr(), n, top_k); +} +} // namespace + +void gguf_dp4a_raw_down_grouped_sm70_out( + torch::Tensor out, torch::Tensor projected, torch::Tensor input, + torch::Tensor ids, torch::Tensor order, torch::Tensor probabilities, + torch::Tensor weights, int64_t source_type, int64_t left) { + const c10::cuda::CUDAGuard guard(input.device()); + require_sm70(); + TORCH_CHECK( + (source_type == 20 || source_type == 42) && left >= 0 && left % 32 == 0, + "Grouped down requires original IQ4_NL or Q2_0 rows"); + TORCH_CHECK(input.is_cuda() && input.scalar_type() == at::kByte && + input.is_contiguous() && input.dim() == 4 && + input.size(0) > 0 && input.size(0) <= 40 && + input.size(1) > 0 && input.size(1) <= 16 && + input.size(2) > 0 && input.size(3) == sizeof(Q8_1), + "Invalid grouped down Q8_1 inputs"); + const int m = input.size(0), top_k = input.size(1), k = input.size(2) * 32; + TORCH_CHECK(ids.device() == input.device() && ids.scalar_type() == at::kInt && + ids.is_contiguous() && ids.dim() == 2 && ids.size(0) == m && + ids.size(1) == top_k && order.device() == input.device() && + order.scalar_type() == at::kInt && order.is_contiguous() && + order.sizes() == ids.sizes() && + probabilities.device() == input.device() && + probabilities.scalar_type() == at::kFloat && + probabilities.is_contiguous() && + probabilities.sizes() == ids.sizes(), + "Invalid grouped down route mapping"); + TORCH_CHECK( + out.device() == input.device() && out.scalar_type() == at::kHalf && + out.is_contiguous() && out.dim() == 2 && out.size(0) == m && + out.size(1) > 0 && weights.device() == input.device() && + weights.scalar_type() == at::kByte && weights.is_contiguous() && + weights.dim() == 3 && weights.size(0) > 0 && + weights.size(1) == out.size(1) && + weights.size(2) >= ((k + left + (source_type == 20 ? 31 : 63)) / + (source_type == 20 ? 32 : 64)) * + 18 && + projected.device() == input.device() && + projected.scalar_type() == at::kHalf && projected.is_contiguous() && + projected.numel() == int64_t(m) * top_k * out.size(1), + "Invalid grouped down banks or output"); + if (source_type == 20) + launch_grouped_raw_down<20>(out, projected, input, ids, order, + probabilities, weights, left); + else + launch_grouped_raw_down<42>(out, projected, input, ids, order, + probabilities, weights, left); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh new file mode 100644 index 0000000000..18211ad064 --- /dev/null +++ b/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Sorted route groups reuse original GGUF weight/codebook decoding. Integer +// dots, group reduction, FP16 activation boundaries and Q8_1 encoding retain +// the single-route order. The original gate/up operator stays unchanged. +namespace { +template +__global__ __launch_bounds__(512, 1) void gate_up_grouped( + Q8_1* output, const Q8_1* activation, const int* ids, + const int* route_tokens, const int* order, const int* original_ids, + const int* mode, const uint8_t* gate, const uint8_t* up, int m, int n, + int k, int stride, int top_k, int experts) { + using Dot = LatticeDot; + __shared__ uint32_t book[Dot::kBookWords], masks[16]; + __shared__ half intermediate[4][32]; + __shared__ int active, count, tokens[4], expert; + extern __shared__ Q8_1 shared_x[]; + const int route = blockIdx.y, routes = m * top_k; + if (mode[0] == 0) return; + if (threadIdx.x == 0) { + const int e = ids[route]; + int first = route; + while (first > 0 && ids[first - 1] == e) --first; + active = ((route - first) % 4 == 0) && e >= 0 && e < experts; + count = 1; + while (count < 4 && route + count < routes && ids[route + count] == e) + ++count; + for (int q = 0; q < count; ++q) + tokens[q] = min(max(route_tokens[route + q], 0), m - 1); + expert = e; + } + __syncthreads(); + if (!active) return; + const int groups = k / 32; + for (int i = threadIdx.x; i < count * groups * 9; i += blockDim.x) { + const int token = tokens[i / (groups * 9)]; + reinterpret_cast(shared_x)[i] = + reinterpret_cast( + activation)[int64_t(token) * groups * 9 + i % (groups * 9)]; + } + Dot::initialize(book, masks); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int local_row = warp * 2 + lane / 16, row = blockIdx.x * 32 + local_row; + const uint8_t* g = gate + (int64_t(expert) * n + row) * stride; + const uint8_t* u = up + (int64_t(expert) * n + row) * stride; + float gs[4] = {}, us[4] = {}; + for (int group = lane % 16; group < groups; group += 16) { + if (count == 1) { + gs[0] += Dot::dot(g, group, shared_x[group], book, masks); + us[0] += Dot::dot(u, group, shared_x[group], book, masks); + } else { + const auto gd = Dot::decode(g, group, book, masks); + const auto ud = Dot::decode(u, group, book, masks); +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < count) { + gs[q] += Dot::dot_decoded(gd, shared_x[q * groups + group]); + us[q] += Dot::dot_decoded(ud, shared_x[q * groups + group]); + } + } + } + } +#pragma unroll + for (int offset = 8; offset; offset >>= 1) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + gs[q] += __shfl_down_sync(0xffffffff, gs[q], offset, 16); + us[q] += __shfl_down_sync(0xffffffff, us[q], offset, 16); + } + } + if (lane % 16 == 0) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < count) { + const float gv = __half2float(__float2half_rn(gs[q])); + intermediate[q][local_row] = __hmul( + __float2half_rn(gv / (1.f + expf(-gv))), __float2half_rn(us[q])); + } + } + } + __syncthreads(); + if (threadIdx.x < 32) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < count) + vllm::sm70_gguf::quantize_q8_1_warp( + output + int64_t(order[route + q]) * (n / 32) + blockIdx.x, + __half2float(intermediate[q][threadIdx.x])); + } + } +} + +template +void launch_grouped_gate(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, torch::Tensor tokens, + torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + gate_up_grouped<<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(activation.data_ptr()), ids.data_ptr(), + tokens.data_ptr(), order.data_ptr(), + original_ids.data_ptr(), mode.data_ptr(), + gate.data_ptr(), up.data_ptr(), activation.size(0), + gate.size(1), activation.size(1) * 32, gate.size(2), ids.size(1), + gate.size(0)); + // Sparse routing uses the retained single-route kernel, with only a device + // admission guard. It must not inherit the grouped kernel's four-row register + // and shared-memory footprint. Both guards are graph replay inputs. + gate_up + <<>>( + out.data_ptr(), reinterpret_cast(activation.data_ptr()), + original_ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), gate.size(1), activation.size(1) * 32, + gate.size(2), ids.size(1), nullptr, nullptr, mode.data_ptr()); +} +} // namespace + +template +static void grouped_gate_out_impl(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, int64_t source_type, + bool bank_aware, torch::Tensor route_tokens, + torch::Tensor order, + torch::Tensor original_ids, + torch::Tensor mode) { + TORCH_CHECK(source_type == 18 || source_type == 21 || source_type == 22, + "Paired gate/up requires an original lattice format"); + const c10::cuda::CUDAGuard guard(activation.device()); + require_sm70(); + TORCH_CHECK(activation.is_cuda() && activation.scalar_type() == at::kByte && + activation.is_contiguous() && activation.dim() == 3 && + activation.size(0) > 0 && activation.size(0) <= MaxRows && + activation.size(2) == sizeof(Q8_1) && + activation.size(1) % 8 == 0, + "Invalid paired Q8_1 activation"); + const int k = activation.size(1) * 32; + const int bytes = source_type == 18 ? 98 : source_type == 21 ? 110 : 82; + TORCH_CHECK(gate.device() == activation.device() && + up.device() == activation.device() && + gate.scalar_type() == at::kByte && + up.scalar_type() == at::kByte && gate.is_contiguous() && + up.is_contiguous() && gate.dim() == 3 && + gate.sizes() == up.sizes() && gate.size(0) > 0 && + gate.size(1) > 0 && gate.size(1) % 32 == 0 && + (gate.size(2) == k / 256 * bytes || + gate.size(2) == (k / 256 * bytes + 7) / 8 * 8), + "Invalid paired original weight banks"); + TORCH_CHECK( + ids.device() == activation.device() && ids.scalar_type() == at::kInt && + ids.is_contiguous() && ids.dim() == 2 && + ids.size(0) == activation.size(0) && ids.size(1) > 0 && + ids.size(1) <= 16 && route_tokens.device() == activation.device() && + route_tokens.scalar_type() == at::kInt && + route_tokens.is_contiguous() && route_tokens.numel() == ids.numel(), + "Invalid sorted route or token mapping"); + TORCH_CHECK(order.device() == activation.device() && + order.scalar_type() == at::kInt && order.is_contiguous() && + order.numel() == ids.numel() && + original_ids.device() == activation.device() && + original_ids.scalar_type() == at::kInt && + original_ids.is_contiguous() && + original_ids.sizes() == ids.sizes() && + mode.device() == activation.device() && + mode.scalar_type() == at::kInt && mode.is_contiguous() && + mode.numel() == 1, + "Invalid grouped routing workspace"); + TORCH_CHECK( + out.device() == activation.device() && out.scalar_type() == at::kByte && + out.is_contiguous() && out.dim() == 4 && + out.size(0) == activation.size(0) && out.size(1) == ids.size(1) && + out.size(2) == gate.size(1) / 32 && out.size(3) == sizeof(Q8_1), + "Invalid paired Q8_1 output"); +#define PAIR(TYPE) \ + if (bank_aware) \ + launch_grouped_gate(out, activation, ids, gate, up, \ + route_tokens, order, original_ids, mode); \ + else \ + launch_grouped_gate(out, activation, ids, gate, up, \ + route_tokens, order, original_ids, mode) + if (source_type == 18) { + PAIR(18); + } else if (source_type == 21) { + PAIR(21); + } else { + PAIR(22); + } +#undef PAIR + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void gguf_dp4a_gate_up_grouped_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + grouped_gate_out_impl<20>(out, activation, ids, gate, up, source_type, + bank_aware, route_tokens, order, original_ids, + mode); +} + +void gguf_dp4a_gate_up_grouped_batch_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + grouped_gate_out_impl<40>(out, activation, ids, gate, up, source_type, + bank_aware, route_tokens, order, original_ids, + mode); +} diff --git a/csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh b/csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh new file mode 100644 index 0000000000..ddb721e464 --- /dev/null +++ b/csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh @@ -0,0 +1,378 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Included in namespace hcx, after the retained HCX primitives. Full-mesh +// TP4, separate output projection, local schedule, M9..40 only. Sixteen +// warps expose independent row exchange/reduction work; the five HC up warps +// and six HC down warps retain their original arithmetic. One resident +// grid performs the exchange and both grid rendezvous for the whole batch. +constexpr int BATCH_ROWS = 40; +constexpr int BATCH_THREADS = 512; + +template +__global__ __launch_bounds__(BATCH_THREADS, 1) void hcx_batch_kernel(Args a) { + constexpr int WARPS = BATCH_THREADS / 32; + const int i = blockIdx.x, t = threadIdx.x, warp = t >> 5, lane = t & 31; + const int rk = a.rank, p1 = rk ^ 1, p2 = rk ^ 2, p3 = rk ^ 3; + const unsigned tag = ldv(a.seq) + 1, base = ldv(a.bar + 3); + const int r8 = (lane & 3) + ((lane & 16) ? 4 : 0), quad = (lane >> 2) & 3; + constexpr int AR_STRIDE = BATCH_ROWS * 32; + __shared__ __align__(16) half xs[BATCH_ROWS][136]; + __shared__ __align__(16) half ls[BATCH_ROWS][328]; + // The LoRA input tile is not populated until after the down projection + // and both grid rendezvous. Reuse its earlier lifetime for a coalesced + // split-K partial tile instead of allocating another shared-memory bank. + static_assert(sizeof(ls) >= 8 * 388 * sizeof(float)); + float* const partial_tile = reinterpret_cast(ls); + __shared__ float red[5][8][33]; + __shared__ float rr[BATCH_ROWS][4]; + + if constexpr (HalfPackets) { + for (int row = warp; row < a.M; row += WARPS) { + const int col = 32 * i + lane; + const bool packed = a.p1 == nullptr; + const int slot = (rk * NC + i) * AR_STRIDE + row * 32 + lane; + if (!packed || lane < 16) { + unsigned value; + if (packed) { + const half* pair = a.p0 + row * HD + 32 * i + 2 * lane; + value = static_cast(__half_as_ushort(pair[0])) | + (static_cast(__half_as_ushort(pair[1])) << 16); + } else { + float v = __half2float(a.p0[row * HD + col]); + v += __half2float(a.p1[row * HD + col]); + value = __float_as_uint(v); + } + st2(a.ar[p1] + slot, value, tag); + st2(a.ar[p2] + slot, value, tag); + st2(a.ar[p3] + slot, value, tag); + } + // A tagged header in previously unused LoRA padding preserves each peer's + // optional-secondary semantics, including mixed formats across TP ranks. + // Every data word retains its existing payload+epoch validation. + if (i == 0 && lane == 0) { + const int header = row * LW + 324 + rk; + st2(a.lora[p1] + header, packed, tag); + st2(a.lora[p2] + header, packed, tag); + st2(a.lora[p3] + header, packed, tag); + } + } + } else { + for (int row = warp; row < a.M; row += WARPS) { + const int col = 32 * i + lane; + float v = __half2float(a.p0[row * HD + col]); + if (a.p1) v += __half2float(a.p1[row * HD + col]); + const int slot = (rk * NC + i) * AR_STRIDE + row * 32 + lane; + st2(a.ar[p1] + slot, __float_as_uint(v), tag); + st2(a.ar[p2] + slot, __float_as_uint(v), tag); + st2(a.ar[p3] + slot, __float_as_uint(v), tag); + } + } + uint4 dlo[4], dhi[4], ulo[4], uhi[4]; + if (warp < 6) { +#pragma unroll + for (int s = 0; s < 4; ++s) { + const uint4* p = a.wd + ((i * 6 + warp) * 4 + s) * 64; + dlo[s] = ldgv(p + lane); + dhi[s] = ldgv(p + 32 + lane); + } + } + if (warp >= 3 && warp < 8) { +#pragma unroll + for (int s = 0; s < 4; ++s) { + const uint4* p = a.wu + ((i * 5 + warp - 3) * 4 + s) * 64; + ulo[s] = ldgv(p + lane); + uhi[s] = ldgv(p + 32 + lane); + } + } + for (int row = warp; row < BATCH_ROWS; row += WARPS) { + const int col = 32 * i + lane; + float sums[4] = {}; +#pragma unroll + for (int b = 0; b < 4; ++b) xs[row][b * 32 + lane] = __float2half_rn(0.f); + if (row < a.M) { + float v = __half2float(a.p0[row * HD + col]); + if (a.p1) v += __half2float(a.p1[row * HD + col]); + unsigned f1 = 0, f2 = 0, f3 = 0; + if constexpr (HalfPackets) { + f1 = poll2(a.lora[rk] + row * LW + 324 + p1, tag); + f2 = poll2(a.lora[rk] + row * LW + 324 + p2, tag); + f3 = poll2(a.lora[rk] + row * LW + 324 + p3, tag); + } + const int slot = i * AR_STRIDE + row * 32; + // Retain the qualified small-HC schedule: issue all independent peer + // loads before waiting for any one epoch. Arithmetic order is unchanged. + unsigned u1, u2, u3, e1, e2, e3; + const auto start = clock64(); + do { + asm volatile("ld.volatile.global.v2.u32 {%0,%1}, [%2];" + : "=r"(u1), "=r"(e1) + : "l"(a.ar[rk] + p1 * NC * AR_STRIDE + slot + + (f1 ? lane / 2 : lane))); + asm volatile("ld.volatile.global.v2.u32 {%0,%1}, [%2];" + : "=r"(u2), "=r"(e2) + : "l"(a.ar[rk] + p2 * NC * AR_STRIDE + slot + + (f2 ? lane / 2 : lane))); + asm volatile("ld.volatile.global.v2.u32 {%0,%1}, [%2];" + : "=r"(u3), "=r"(e3) + : "l"(a.ar[rk] + p3 * NC * AR_STRIDE + slot + + (f3 ? lane / 2 : lane))); + if (clock64() - start > 4000000000LL) __trap(); + } while (e1 != tag || e2 != tag || e3 != tag); + const int shift = 16 * (lane & 1); + const float x1 = f1 ? __half2float(__ushort_as_half(u1 >> shift)) + : __uint_as_float(u1); + const float x2 = f2 ? __half2float(__ushort_as_half(u2 >> shift)) + : __uint_as_float(u2); + const float x3 = f3 ? __half2float(__ushort_as_half(u3 >> shift)) + : __uint_as_float(u3); + const float near = (rk & 1) ? x1 + v : v + x1; + const float far = (rk & 1) ? x3 + x2 : x2 + x3; + const float reduced = + __half2float(__float2half_rn((rk & 2) ? far + near : near + far)); +#pragma unroll + for (int b = 0; b < 4; ++b) { + const float g = 2.f * sigm(__half2float(a.inj[row * 4 + b]) / 4.f); + const half h = __float2half_rn( + __half2float(a.res[row * KD + b * HD + col]) + reduced * g); + const float o = __half2float(h); + const half nw = a.nw[a.nw_full ? b * HD + col : col]; + a.res_out[row * KD + b * HD + col] = h; + xs[row][b * 32 + lane] = __float2half_rn(o + o * __half2float(nw)); + sums[b] = o * o; + } + } +#pragma unroll + for (int s = 16; s > 0; s >>= 1) +#pragma unroll + for (int b = 0; b < 4; ++b) + sums[b] += __shfl_xor_sync(0xffffffff, sums[b], s); + if (lane == 0 && row < a.M) + *reinterpret_cast(a.sq + (i * BATCH_ROWS + row) * 4) = + make_float4(sums[0], sums[1], sums[2], sums[3]); + } + __syncthreads(); + const int tile = warp >> 1, kh = warp & 1; +#pragma unroll + for (int token_tile = 0; token_tile < 5; ++token_tile) { + if (token_tile * 8 >= a.M) break; + if (warp < 6) { + float acc0[8] = {}, acc1[8] = {}; +#pragma unroll + for (int s = 0; s < 4; ++s) { + const int k = (kh * 4 + s) * 16; + const uint4 x0 = + *reinterpret_cast(&xs[token_tile * 8 + r8][k]); + const uint4 x1 = + *reinterpret_cast(&xs[token_tile * 8 + r8][k + 8]); + float (&acc)[8] = s < 2 ? acc0 : acc1; + mma(acc, x0.x, x0.y, dlo[s].x, dlo[s].y); + mma(acc, x0.z, x0.w, dlo[s].z, dlo[s].w); + mma(acc, x1.x, x1.y, dhi[s].x, dhi[s].y); + mma(acc, x1.z, x1.w, dhi[s].z, dhi[s].w); + } +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int rw = + token_tile * 8 + (e & 2) + ((lane & 16) ? 4 : 0) + (lane & 1); + const int n = tile * 32 + quad * 8 + + ((e & 1) | (((lane >> 1) & 1) << 1) | ((e >> 2) << 2)); + *reinterpret_cast(partial_tile + (rw - token_tile * 8) * 388 + + n * 4 + 2 * kh) = + make_float2(acc0[e], acc1[e]); + } + } + __syncthreads(); + for (int cell = t; cell < 8 * 96; cell += BATCH_THREADS) { + const int row = cell / 96, col = cell % 96; + if (token_tile * 8 + row < a.M) + *reinterpret_cast( + a.dpart + + ((i * BATCH_ROWS + token_tile * 8 + row) * 96 + col) * 4) = + *reinterpret_cast(partial_tile + row * 388 + + col * 4); + } + __syncthreads(); + } + // Reserved padded down columns84..95 have no downstream readers. The + // producing CTA resets each row flag after all its partial stores, before + // the existing release/acquire grid barrier. No allocation/init dependency. + if (i == NC - 1 && t < a.M) { + float* slot = a.dpart + ((NC - 1) * BATCH_ROWS + t) * 96 * 4 + 84 * 4; + *reinterpret_cast(slot + 4) = tag - 1; + } + grid_bar_epoch(a.bar + 2, base + NC); + // A resident owner computes the identical80-way statistic once per row. + if (i < a.M && warp == 0) { + const int row = i; + float4 qs[3]; +#pragma unroll + for (int j = 0; j < 3; ++j) + qs[j] = lane + 32 * j < NC + ? __ldcg(reinterpret_cast(a.sq) + + (lane + 32 * j) * BATCH_ROWS + row) + : make_float4(0.f, 0.f, 0.f, 0.f); + float sums[4] = { + (qs[0].x + qs[1].x) + qs[2].x, (qs[0].y + qs[1].y) + qs[2].y, + (qs[0].z + qs[1].z) + qs[2].z, (qs[0].w + qs[1].w) + qs[2].w}; +#pragma unroll + for (int s = 16; s > 0; s >>= 1) +#pragma unroll + for (int b = 0; b < 4; ++b) + sums[b] += __shfl_xor_sync(0xffffffff, sums[b], s); + + if (lane == 0) { + float* slot = a.dpart + ((NC - 1) * BATCH_ROWS + row) * 96 * 4 + 84 * 4; + *reinterpret_cast(slot) = make_float4( + rsqrtf(sums[0] / HD + a.eps), rsqrtf(sums[1] / HD + a.eps), + rsqrtf(sums[2] / HD + a.eps), rsqrtf(sums[3] / HD + a.eps)); + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(slot + 4), + "r"(tag) + : "memory"); + } + } + for (int row = warp; row < a.M; row += WARPS) { + float inverse_values[4] = {}; + if (lane == 0) { + const float* slot = + a.dpart + ((NC - 1) * BATCH_ROWS + row) * 96 * 4 + 84 * 4; + const auto start = clock64(); + unsigned ready; + do { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" + : "=r"(ready) + : "l"(slot + 4) + : "memory"); + if (clock64() - start > 4000000000LL) __trap(); + } while (ready != tag); + const float4 value = __ldcg(reinterpret_cast(slot)); + inverse_values[0] = value.x; + inverse_values[1] = value.y; + inverse_values[2] = value.z; + inverse_values[3] = value.w; + } +#pragma unroll + for (int b = 0; b < 4; ++b) { + const float inverse = __shfl_sync(0xffffffff, inverse_values[b], 0); + if (lane == 0) rr[row][b] = inverse; + const float o = + __half2float(a.res_out[row * KD + b * HD + 32 * i + lane]); + float y = o * inverse; + y += y * __half2float( + a.nw[a.nw_full ? b * HD + 32 * i + lane : 32 * i + lane]); + a.xn[row * KD + b * HD + 32 * i + lane] = __float2half_rn(y); + } + } + __syncthreads(); + for (int oi = i + NC * warp; oi < a.M * 96; oi += NC * WARPS) { + const int rw = oi / 96, n = oi % 96; + if (!(n < 80 || (rk == 3 && n < 84))) continue; + float pb[4] = {}; +#pragma unroll + for (int j = 0; j < 3; ++j) { + if (lane + 32 * j < NC) { + const float4 pv = __ldcg(reinterpret_cast(a.dpart) + + ((lane + 32 * j) * BATCH_ROWS + rw) * 96 + n); + pb[0] += pv.x; + pb[1] += pv.y; + pb[2] += pv.z; + pb[3] += pv.w; + } + } +#pragma unroll + for (int s = 16; s > 0; s >>= 1) +#pragma unroll + for (int b = 0; b < 4; ++b) + pb[b] += __shfl_xor_sync(0xffffffff, pb[b], s); + float q = 0.f; +#pragma unroll + for (int b = 0; b < 4; ++b) q = fmaf(rr[rw][b], pb[b], q); + half hv = __float2half_rn(q); + if (n < 80) { + const float v = divf(__half2float(hv), 4.f); + hv = __float2half_rn(v * sigm(v)); + } + const int pos = rw * LW + (n < 80 ? rk * 80 + n : 320 + n - 80); + const unsigned value = __half_as_ushort(hv); + if (lane == 0) st2(a.lora[rk] + pos, value, tag); + if (lane == 1) st2(a.lora[p1] + pos, value, tag); + if (lane == 2) st2(a.lora[p2] + pos, value, tag); + if (lane == 3) st2(a.lora[p3] + pos, value, tag); + } + for (int idx = i * BATCH_THREADS + t; idx < a.M * 324; + idx += NC * BATCH_THREADS) + poll2(a.lora[rk] + (idx / 324) * LW + idx % 324, tag); + grid_bar_epoch(a.bar + 2, base + 2 * NC); + for (int idx = t; idx < BATCH_ROWS * 320; idx += BATCH_THREADS) + ls[idx / 320][idx % 320] = + idx < a.M * 320 + ? __ushort_as_half(static_cast( + __ldcg(a.lora[rk] + (idx / 320) * LW + idx % 320).x)) + : __float2half_rn(0.f); + if (i == 0 && t < a.M * 4) + a.inj_out[t] = __ushort_as_half(static_cast( + __ldcg(a.lora[rk] + (t >> 2) * LW + 320 + (t & 3)).x)); + __syncthreads(); +#pragma unroll + for (int token_tile = 0; token_tile < 5; ++token_tile) { + if (token_tile * 8 >= a.M) break; + if (warp >= 3 && warp < 8) { + const int uw = warp - 3; + float acc[8] = {}; +#pragma unroll + for (int s = 0; s < 4; ++s) { + const int k = (uw * 4 + s) * 16; + const uint4 x0 = + *reinterpret_cast(&ls[token_tile * 8 + r8][k]); + const uint4 x1 = + *reinterpret_cast(&ls[token_tile * 8 + r8][k + 8]); + mma(acc, x0.x, x0.y, ulo[s].x, ulo[s].y); + mma(acc, x0.z, x0.w, ulo[s].z, ulo[s].w); + mma(acc, x1.x, x1.y, uhi[s].x, uhi[s].y); + mma(acc, x1.z, x1.w, uhi[s].z, uhi[s].w); + } +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int rw = (e & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + red[uw][rw][quad * 8 + ((e & 1) | (lane & 2) | (e & 4))] = acc[e]; + } + } + __syncthreads(); + const int row = token_tile * 8 + warp; + if (warp < 8 && row < a.M) { + const int h = 640 * rk + 8 * i + (lane & 7); + float acc = red[0][warp][lane]; +#pragma unroll + for (int q = 1; q < 5; ++q) acc += red[q][warp][lane]; + const float sg = sigm(__half2float(__float2half_rn(acc))); + const float xv = + __half2float(__ldcg(a.xn + row * KD + (lane >> 3) * HD + h)); + float mixed = 0.f; +#pragma unroll + for (int b = 0; b < 4; ++b) + mixed = fmaf(__shfl_sync(0xffffffff, sg, b * 8 + (lane & 7)), + __shfl_sync(0xffffffff, xv, b * 8 + (lane & 7)), mixed); + const unsigned bits = __half_as_ushort(__float2half_rn(divf(mixed, 4.f))); + const unsigned packed = + bits | (__shfl_down_sync(0xffffffff, bits, 1) << 16); + const unsigned upper = __shfl_down_sync(0xffffffff, packed, 2); + if (lane < 8 && !(lane & 3)) { + st4(a.hb[p1] + (row * HD + h) / 2, packed, upper, tag); + st4(a.hb[p2] + (row * HD + h) / 2, packed, upper, tag); + st4(a.hb[p3] + (row * HD + h) / 2, packed, upper, tag); + *reinterpret_cast(a.blk_out + row * HD + h) = + make_uint2(packed, upper); + } + } + __syncthreads(); + } + for (int u = t; u < 6 * a.M; u += BATCH_THREADS) { + const int src = u % 3 == 0 ? p1 : u % 3 == 1 ? p2 : p3; + const int row = (u / 3) >> 1, h = 640 * src + 8 * i + ((u / 3) & 1) * 4; + *reinterpret_cast(a.blk_out + row * HD + h) = + poll4(a.hb[rk] + (row * HD + h) / 2, tag); + } + if (i == 0 && t == 0) { + a.bar[3] = base + 2 * NC; + a.seq[0] = tag; + } +} diff --git a/csrc/sm70_turbomind/ops/hcx_sm70.cu b/csrc/sm70_turbomind/ops/hcx_sm70.cu index 8bb2461974..88520e4753 100644 --- a/csrc/sm70_turbomind/ops/hcx_sm70.cu +++ b/csrc/sm70_turbomind/ops/hcx_sm70.cu @@ -736,9 +736,11 @@ __global__ __launch_bounds__(256, 1) void hcx_kernel(Args a) { } } } +#include "hcx_batch_sm70.cuh" } // namespace hcx -void sm70_hcx_out( +template +void sm70_hcx_impl( torch::Tensor p0, std::optional p1, torch::Tensor res, torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, @@ -779,6 +781,23 @@ void sm70_hcx_out( a.dbg = dbg ? reinterpret_cast(dbg->data_ptr()) : nullptr; a.M = static_cast(p0.size(0)); + if constexpr (Batch) { + TORCH_CHECK(a.M >= 9 && a.M <= 40 && full && local_schedule && !ox && !gz, + "HCX batch requires M9..40, full mesh, local schedule and " + "separate projection"); + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0 && + props->multiProcessorCount >= hcx::NC, + "HCX batch requires SM70 with at least 80 SMs"); + TORCH_CHECK(xn.numel() >= 40 * hcx::KD && sq.numel() >= hcx::NC * 40 * 4 && + dpart.numel() >= hcx::NC * 40 * 96 * 4 && bar.numel() >= 4, + "HCX batch workspace too small"); + hcx::hcx_batch_kernel + <<>>( + a); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } TORCH_CHECK(a.M >= 1 && a.M <= 8, "HCX: M1..8"); auto st = at::cuda::getCurrentCUDAStream(); if (ox) { @@ -878,3 +897,60 @@ void sm70_dmv13_out(torch::Tensor x, std::vector codes, dmvns::run(x, codes, high, scale, out, fmt, n, K, split, warps, ws, cnt, tp, std::nullopt, std::nullopt, extra_weight, extra_out); } + +void sm70_hcx_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule) { + sm70_hcx_impl(p0, p1, res, inj, nw, eps, wd, wu, res_out, blk_out, + inj_out, xn, sq, dpart, bar, seq, ar, lora, hb, rank, + dbg, full, ox, ocodes, ohigh, oscale, ofmt, gz, gw, geps, + gscr, local_schedule); +} + +void sm70_hcx_batch_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule) { + sm70_hcx_impl(p0, p1, res, inj, nw, eps, wd, wu, res_out, blk_out, + inj_out, xn, sq, dpart, bar, seq, ar, lora, hb, rank, dbg, + full, ox, ocodes, ohigh, oscale, ofmt, gz, gw, geps, gscr, + local_schedule); +} + +void sm70_hcx_batch_packed_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule) { + sm70_hcx_impl(p0, p1, res, inj, nw, eps, wd, wu, res_out, blk_out, + inj_out, xn, sq, dpart, bar, seq, ar, lora, hb, + rank, dbg, full, ox, ocodes, ohigh, oscale, ofmt, + gz, gw, geps, gscr, local_schedule); +} diff --git a/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu b/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu index d332eb2ce0..42ee3e3a3b 100644 --- a/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu +++ b/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu @@ -23,8 +23,9 @@ struct Args { const int64_t* positions; const int* lengths; const half* gate; - const half* hot; // [hot_tokens, 2, D] FP16, or null - const int* resolved; // [M, width] hot token, -1 host, -2 invalid + const half* hot; // [hot_tokens, 2, D] FP16 or encoded E4M3 bytes + const float* hot_scales; // [hot_tokens, 2], encoded hot storage only + const int* resolved; // [M, width] hot token, -1 host, -2 invalid half* out; float* partial; float* state; @@ -140,14 +141,24 @@ __device__ __forceinline__ uint4 decode8(const unsigned char* source, // Read one 8-value chunk of token K (kv=0) or V (kv=1): from the protected // hot cache when resolved there, else from the authoritative history, which // may be pinned host memory mapped into the device address space (UVA). -template +template __device__ __forceinline__ uint4 load_kv(const Args& a, int64_t slot, int hot, int kv, int dim, const float* lut) { if (slot < 0) return make_uint4(0, 0, 0, 0); if constexpr (HOT) { - if (hot >= 0) - return __ldg(reinterpret_cast( - a.hot + (int64_t(hot) * 2 + kv) * D + dim)); + if (hot >= 0) { + if constexpr (HOT_FP8) { + const auto* source = reinterpret_cast(a.hot) + + (int64_t(hot) * 2 + kv) * D + dim; + // Match the software bit decoder used to materialize FP16 hot values, + // including the FP16 rounding before QK/PV arithmetic. + return decode8_bits(__ldg(reinterpret_cast(source)), + __ldg(a.hot_scales + int64_t(hot) * 2 + kv)); + } else { + return __ldg(reinterpret_cast( + a.hot + (int64_t(hot) * 2 + kv) * D + dim)); + } + } } const int block = slot / a.page, off = slot % a.page; const int64_t element = ((int64_t(block) * 2 + kv) * a.page + off) * D + dim; @@ -158,7 +169,7 @@ __device__ __forceinline__ uint4 load_kv(const Args& a, int64_t slot, int hot, reinterpret_cast(a.history) + element)); } -template +template __global__ __launch_bounds__(256) void partial(Args a) { const int row = blockIdx.y, split = blockIdx.x; const int t = threadIdx.x, warp = t >> 5, lane = t & 31; @@ -211,7 +222,7 @@ __global__ __launch_bounds__(256) void partial(Args a) { for (int v = t; v < C * 32; v += 256) { const int token = v / 32, dim = (v % 32) * 8; *reinterpret_cast(&kv[token][dim]) = - load_kv(a, slots[token], hots[token], 0, dim, lut); + load_kv(a, slots[token], hots[token], 0, dim, lut); } __syncthreads(); // Without the hot cache, issue this thread's V loads now so their @@ -296,8 +307,8 @@ __global__ __launch_bounds__(256) void partial(Args a) { if constexpr (HOT) { for (int v = t; v < C * 32; v += 256) { const int token = v / 32, dim = (v % 32) * 8; - *reinterpret_cast(&kv[token][dim]) = - load_kv(a, slots[token], hots[token], 1, dim, lut); + *reinterpret_cast(&kv[token][dim]) = load_kv( + a, slots[token], hots[token], 1, dim, lut); } } else { #pragma unroll @@ -779,7 +790,8 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, std::optional hot = std::nullopt, std::optional resolved = std::nullopt, int64_t tiles = 1, bool prefill = false, - int64_t whole_rows = 0) { + int64_t whole_rows = 0, + std::optional hot_scales = std::nullopt) { const c10::cuda::CUDAGuard guard(q.device()); TORCH_CHECK(q.is_cuda() && q.scalar_type() == at::kHalf && q.dim() == 3 && q.size(0) > 0 && (whole_rows || q.size(0) <= 64) && @@ -813,16 +825,29 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, TORCH_CHECK(tensor.device() == q.device()); if (gate) TORCH_CHECK(gate->device() == q.device()); TORCH_CHECK(hot.has_value() == resolved.has_value()); + TORCH_CHECK(!hot_scales || (hot && history.scalar_type() == at::kByte), + "Encoded hot storage requires E4M3 authoritative history"); if (hot) { TORCH_CHECK(hot->is_cuda() && hot->device() == q.device() && - hot->scalar_type() == at::kHalf && hot->dim() == 3 && - hot->size(1) == 2 && hot->size(2) == 256 && + hot->scalar_type() == (hot_scales ? at::kByte : at::kHalf) && + hot->dim() == 3 && hot->size(1) == 2 && hot->size(2) == 256 && hot->is_contiguous()); + TORCH_CHECK( + reinterpret_cast(hot->data_ptr()) % (hot_scales ? 8 : 16) == + 0, + "QSA hot storage requires aligned vector loads"); TORCH_CHECK(resolved->device() == q.device() && resolved->scalar_type() == at::kInt && resolved->dim() == 2 && resolved->size(0) >= q.size(0) && resolved->size(1) >= indices.size(1) && resolved->stride(1) == 1); + if (hot_scales) { + TORCH_CHECK(hot_scales->device() == q.device() && + hot_scales->scalar_type() == at::kFloat && + hot_scales->dim() == 2 && + hot_scales->size(0) == hot->size(0) && + hot_scales->size(1) == 2 && hot_scales->is_contiguous()); + } } TORCH_CHECK(reinterpret_cast(q.data_ptr()) % 16 == 0 && q.stride(0) % 8 == 0 && q.stride(1) % 8 == 0); @@ -839,6 +864,7 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, a.lengths = lengths.data_ptr(); a.gate = gate ? reinterpret_cast(gate->data_ptr()) : nullptr; a.hot = hot ? reinterpret_cast(hot->data_ptr()) : nullptr; + a.hot_scales = hot_scales ? hot_scales->data_ptr() : nullptr; a.resolved = resolved ? resolved->data_ptr() : nullptr; a.rs0 = resolved ? resolved->stride(0) : 0; a.out = reinterpret_cast(out.data_ptr()); @@ -896,6 +922,8 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, } if (prefill && fp8 && !hot && a.tiles == 1) device_qsa::partial_fp8_lite<<>>(a); + else if (fp8 && hot_scales) + device_qsa::partial<<>>(a); else if (fp8 && hot) device_qsa::partial<<>>(a); else if (fp8) @@ -985,6 +1013,22 @@ TORCH_LIBRARY(vllm_sm70_qsa_device, m) { positions, lengths, out, gate, partial, state, hot, resolved); }); + m.def( + "run_hot_fp8(Tensor q, Tensor history, Tensor scales, Tensor indices, " + "Tensor table, Tensor requests, Tensor positions, Tensor lengths, " + "Tensor(a!) out, Tensor? gate, Tensor(b!) partial, Tensor(c!) state, " + "Tensor hot, Tensor resolved, Tensor hot_scales) -> ()"); + m.impl("run_hot_fp8", torch::kCUDA, + [](torch::Tensor q, torch::Tensor history, torch::Tensor scales, + torch::Tensor indices, torch::Tensor table, torch::Tensor requests, + torch::Tensor positions, torch::Tensor lengths, torch::Tensor out, + std::optional gate, torch::Tensor partial, + torch::Tensor state, torch::Tensor hot, torch::Tensor resolved, + torch::Tensor hot_scales) { + device_qsa_out(q, history, scales, indices, table, requests, + positions, lengths, out, gate, partial, state, hot, + resolved, 1, false, 0, hot_scales); + }); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} diff --git a/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu b/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu index d0be9cf100..4b7af4957e 100644 --- a/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu +++ b/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu @@ -20,14 +20,29 @@ namespace { "+f"(C[5]), "+f"(C[6]), "+f"(C[7]) \ : "r"(A0), "r"(A1), "r"(B0), "r"(B1)) -template -__global__ void mqa_multiquery(const half* q, const half* cache, - const int* table, const int* requests, - const Position* positions, const int* lengths, - float* scores, int* visible, int m, int columns, - int pages, int page_size, int table_width, - int sq0, int sq1, int sq2, int sc0, int sc1, - int sc3, int st1, float divisor, int compress) { +template +__global__ void mqa_multiquery( + const half* q, const half* cache, const int* table, const int* requests, + const Position* positions, const int* lengths, float* scores, int* visible, + int m, int columns, int pages, int page_size, int table_width, int sq0, + int sq1, int sq2, int sc0, int sc1, int sc3, int st1, int st0, + float divisor, int compress, const half* staged, const int* page_slots, + int staged_pages, const half* spill, const int* spill_slots, + int spill_pages) { + // Each grid-Y group keeps the retained M5 MMA/reduction order. Live global + // request IDs also preserve invalid/padded rows without a per-group subtract. + const int request_index = Batch ? blockIdx.y : 0; + if constexpr (Batch) { + q += int64_t(request_index) * 5 * sq0; + table += int64_t(request_index) * st0; + requests += request_index * 5; + positions += request_index * 5; + lengths += request_index; + scores += int64_t(request_index) * 5 * columns; + visible += request_index * 5; + m = 5; + } const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; const int quad = (lane >> 2) & 3; const int r = (lane & 3) + ((lane & 16) ? 4 : 0); @@ -36,7 +51,7 @@ __global__ void mqa_multiquery(const half* q, const half* cache, int end = 0; for (int row = 0; row < m; ++row) { const long long prefix = positions[row] + 1; - const long long length = requests[row] == 0 ? lengths[0] : 0; + const long long length = requests[row] == request_index ? lengths[0] : 0; const long long causal = prefix < length ? prefix : length; // Triton's signed integer division truncates toward zero, including // negative padded positions. Keep its visible-length metadata unchanged. @@ -49,6 +64,30 @@ __global__ void mqa_multiquery(const half* q, const half* cache, ? table[(col / page_size) * st1] : -1; const bool key_valid = col < end && page >= 0 && page < pages; + const half* key = nullptr; + int key_stride = sc3; + if (key_valid) { + int resident = -1; + if constexpr (Staged) resident = page_slots[page]; + if (resident >= 0 && resident < staged_pages) { + key = + staged + + (static_cast(resident) * page_size + col % page_size) * 128; + key_stride = 1; + } else { + int secondary = -1; + if constexpr (Banked) secondary = spill_slots[page]; + if (secondary >= 0 && secondary < spill_pages) { + key = spill + + (static_cast(secondary) * page_size + col % page_size) * + 128; + key_stride = 1; + } else { + key = + cache + static_cast(page) * sc0 + (col % page_size) * sc1; + } + } + } const int qhead = quad * 8 + r, qrow = qhead / 4; float acc[8] = {}; #pragma unroll @@ -59,14 +98,13 @@ __global__ void mqa_multiquery(const half* q, const half* cache, half* bv = reinterpret_cast(&b); half* lv = reinterpret_cast(&lo); half* hv = reinterpret_cast(&hi); - if (sc3 == 1 && sq2 == 1 && sc0 % 8 == 0 && sc1 % 8 == 0 && sq0 % 8 == 0 && - sq1 % 8 == 0 && + if (key_stride == 1 && sq2 == 1 && sc0 % 8 == 0 && sc1 % 8 == 0 && + sq0 % 8 == 0 && sq1 % 8 == 0 && (reinterpret_cast(q) | reinterpret_cast(cache)) % 16 == 0) { if (key_valid) { - const half* k = cache + static_cast(page) * sc0 + - (col % page_size) * sc1 + g * 16; + const half* k = key + g * 16; a = *reinterpret_cast(k); b = *reinterpret_cast(k + 8); } @@ -79,10 +117,8 @@ __global__ void mqa_multiquery(const half* q, const half* cache, #pragma unroll for (int j = 0; j < 8; ++j) { if (key_valid) { - const half* k = cache + static_cast(page) * sc0 + - (col % page_size) * sc1; - av[j] = k[(g * 16 + j) * sc3]; - bv[j] = k[(g * 16 + 8 + j) * sc3]; + av[j] = key[(g * 16 + j) * key_stride]; + bv[j] = key[(g * 16 + 8 + j) * key_stride]; } if (qrow < m) { const half* query = q + qrow * sq0 + (qhead % 4) * sq1; @@ -112,7 +148,7 @@ __global__ void mqa_multiquery(const half* q, const half* cache, const int key_lane = (lane & ~3) | (k * 2 + (lane & 1)); const bool valid = __shfl_sync(0xffffffff, static_cast(key_valid), key_lane); - if ((lane & 2) == 0 && row < m && requests[row] == 0 && + if ((lane & 2) == 0 && row < m && requests[row] == request_index && column < columns && column < (positions[row] + 1) / compress && column < lengths[0] / compress) { scores[row * columns + column] = valid ? sum / divisor : -CUDART_INF_F; @@ -123,15 +159,26 @@ __global__ void mqa_multiquery(const half* q, const half* cache, #undef QSA_MMA -void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, - torch::Tensor cache, torch::Tensor table, torch::Tensor requests, - torch::Tensor positions, torch::Tensor lengths, double divisor, - int64_t compress) { +template +void mqa_out_impl(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, + torch::Tensor cache, torch::Tensor table, + torch::Tensor requests, torch::Tensor positions, + torch::Tensor lengths, double divisor, int64_t compress, + const torch::Tensor& staged, const torch::Tensor& page_slots, + const torch::Tensor& spill = {}, + const torch::Tensor& spill_slots = {}) { const c10::cuda::CUDAGuard guard(q.device()); TORCH_CHECK(q.is_cuda() && q.scalar_type() == at::kHalf && q.dim() == 3 && - q.size(0) > 0 && q.size(0) <= 8 && q.size(1) == 4 && - q.size(2) == 128 && table.dim() == 2 && table.size(0) == 1, - "Shared-key MQA requires FP16 M1..8/H4/D128 and one request"); + q.size(0) > 0 && q.size(0) <= (Batch ? 40 : 8) && + q.size(1) == 4 && q.size(2) == 128 && table.dim() == 2, + "Shared-key MQA requires FP16 queries with H4/D128"); + if constexpr (Batch) { + TORCH_CHECK(table.size(0) >= 1 && table.size(0) <= 8 && + q.size(0) == 5 * table.size(0), + "Batched shared-key MQA requires five rows per request, C1..8"); + } else { + TORCH_CHECK(table.size(0) == 1, "Shared-key MQA requires one request"); + } const int m = q.size(0), columns = scores.size(1); TORCH_CHECK(cache.dim() == 4 && cache.size(2) == 1 && cache.size(3) == 128 && cache.scalar_type() == at::kHalf && @@ -144,7 +191,8 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, "MQA metadata must use same-device int32"); TORCH_CHECK(scores.device() == q.device() && cache.device() == q.device() && requests.numel() >= m && positions.numel() >= m && - visible.numel() >= m && lengths.numel() >= 1, + visible.numel() >= m && + lengths.numel() >= (Batch ? table.size(0) : 1), "MQA tensor extent mismatch"); TORCH_CHECK(positions.device() == q.device() && positions.is_contiguous() && (positions.scalar_type() == at::kLong || @@ -155,10 +203,43 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, "MQA row metadata must be contiguous"); const auto* props = at::cuda::getCurrentDeviceProperties(); TORCH_CHECK(props->major == 7 && props->minor == 0, "SM70 required"); + const half* staged_ptr = nullptr; + const int* slots_ptr = nullptr; + int staged_pages = 0; + if constexpr (Staged) { + TORCH_CHECK( + staged.device() == q.device() && staged.scalar_type() == at::kHalf && + staged.dim() == 4 && staged.is_contiguous() && + staged.size(1) == cache.size(1) && staged.size(2) == 1 && + staged.size(3) == 128 && page_slots.device() == q.device() && + page_slots.scalar_type() == at::kInt && + page_slots.is_contiguous() && page_slots.numel() >= cache.size(0), + "Invalid staged indexer page workspace"); + staged_ptr = reinterpret_cast(staged.data_ptr()); + slots_ptr = page_slots.data_ptr(); + staged_pages = staged.size(0); + } + const half* spill_ptr = nullptr; + const int* spill_slots_ptr = nullptr; + int spill_pages = 0; + if constexpr (Banked) { + TORCH_CHECK( + spill.device() == q.device() && spill.scalar_type() == at::kHalf && + spill.dim() == 4 && spill.is_contiguous() && + spill.size(1) == cache.size(1) && spill.size(2) == 1 && + spill.size(3) == 128 && spill_slots.device() == q.device() && + spill_slots.scalar_type() == at::kInt && + spill_slots.is_contiguous() && spill_slots.numel() >= cache.size(0), + "Invalid spill page workspace"); + spill_ptr = reinterpret_cast(spill.data_ptr()); + spill_slots_ptr = spill_slots.data_ptr(); + spill_pages = spill.size(0); + } if (!columns) return; if (positions.scalar_type() == at::kLong) { - mqa_multiquery - <<<(columns + 31) / 32, 128, 0, at::cuda::getCurrentCUDAStream()>>>( + mqa_multiquery + <<>>( reinterpret_cast(q.data_ptr()), reinterpret_cast(cache.data_ptr()), table.data_ptr(), requests.data_ptr(), @@ -166,10 +247,13 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, scores.data_ptr(), visible.data_ptr(), m, columns, cache.size(0), cache.size(1), table.size(1), q.stride(0), q.stride(1), q.stride(2), cache.stride(0), cache.stride(1), - cache.stride(3), table.stride(1), divisor, compress); + cache.stride(3), table.stride(1), table.stride(0), divisor, + compress, staged_ptr, slots_ptr, staged_pages, spill_ptr, + spill_slots_ptr, spill_pages); } else { - mqa_multiquery - <<<(columns + 31) / 32, 128, 0, at::cuda::getCurrentCUDAStream()>>>( + mqa_multiquery + <<>>( reinterpret_cast(q.data_ptr()), reinterpret_cast(cache.data_ptr()), table.data_ptr(), requests.data_ptr(), @@ -177,10 +261,50 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, scores.data_ptr(), visible.data_ptr(), m, columns, cache.size(0), cache.size(1), table.size(1), q.stride(0), q.stride(1), q.stride(2), cache.stride(0), cache.stride(1), - cache.stride(3), table.stride(1), divisor, compress); + cache.stride(3), table.stride(1), table.stride(0), divisor, + compress, staged_ptr, slots_ptr, staged_pages, spill_ptr, + spill_slots_ptr, spill_pages); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } + +void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, + torch::Tensor cache, torch::Tensor table, torch::Tensor requests, + torch::Tensor positions, torch::Tensor lengths, double divisor, + int64_t compress) { + mqa_out_impl(scores, visible, q, cache, table, requests, positions, + lengths, divisor, compress, {}, {}); +} + +void mqa_staged_out(torch::Tensor scores, torch::Tensor visible, + torch::Tensor q, torch::Tensor cache, torch::Tensor table, + torch::Tensor requests, torch::Tensor positions, + torch::Tensor lengths, double divisor, int64_t compress, + torch::Tensor staged, torch::Tensor page_slots) { + mqa_out_impl(scores, visible, q, cache, table, requests, positions, + lengths, divisor, compress, staged, page_slots); +} +void mqa_banked_out(torch::Tensor scores, torch::Tensor visible, + torch::Tensor q, torch::Tensor cache, torch::Tensor table, + torch::Tensor requests, torch::Tensor positions, + torch::Tensor lengths, double divisor, int64_t compress, + torch::Tensor staged, torch::Tensor page_slots, + torch::Tensor spill, torch::Tensor spill_slots) { + mqa_out_impl(scores, visible, q, cache, table, requests, + positions, lengths, divisor, compress, staged, + page_slots, spill, spill_slots); +} +void mqa_banked_batch_out(torch::Tensor scores, torch::Tensor visible, + torch::Tensor q, torch::Tensor cache, + torch::Tensor table, torch::Tensor requests, + torch::Tensor positions, torch::Tensor lengths, + double divisor, int64_t compress, + torch::Tensor staged, torch::Tensor page_slots, + torch::Tensor spill, torch::Tensor spill_slots) { + mqa_out_impl(scores, visible, q, cache, table, requests, + positions, lengths, divisor, compress, staged, + page_slots, spill, spill_slots); +} } // namespace TORCH_LIBRARY(vllm_sm70_qsa_indexer, m) { m.def( @@ -188,5 +312,24 @@ TORCH_LIBRARY(vllm_sm70_qsa_indexer, m) { "Tensor table, Tensor requests, Tensor positions, Tensor lengths, " "float divisor, int compress) -> ()"); m.impl("run", torch::kCUDA, &mqa_out); + m.def( + "run_staged(Tensor(a!) scores, Tensor(b!) visible, Tensor q, Tensor " + "cache, " + "Tensor table, Tensor requests, Tensor positions, Tensor lengths, " + "float divisor, int compress, Tensor staged, Tensor page_slots) -> ()"); + m.impl("run_staged", torch::kCUDA, &mqa_staged_out); + m.def( + "run_staged_banked(Tensor(a!) scores, Tensor(b!) visible, Tensor q, " + "Tensor cache, Tensor table, Tensor requests, Tensor positions, Tensor " + "lengths, float divisor, int compress, Tensor staged, Tensor page_slots, " + "Tensor spill, Tensor spill_slots) -> ()"); + m.impl("run_staged_banked", torch::kCUDA, &mqa_banked_out); + m.def( + "run_staged_banked_batch(Tensor(a!) scores, Tensor(b!) visible, Tensor " + "q, " + "Tensor cache, Tensor table, Tensor requests, Tensor positions, Tensor " + "lengths, float divisor, int compress, Tensor staged, Tensor page_slots, " + "Tensor spill, Tensor spill_slots) -> ()"); + m.impl("run_staged_banked_batch", torch::kCUDA, &mqa_banked_batch_out); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 32783ee3f0..9ed3875ee3 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -102,6 +102,28 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ", Tensor? oscale, int ofmt, Tensor? gz, Tensor? gw, float geps, Tensor" "(i!)? gscr, bool local_schedule=False) -> ()"); ops.impl("sm70_hcx_out", torch::kCUDA, &sm70_hcx_out); + ops.def( + "sm70_hcx_batch_out(Tensor p0, Tensor? p1, Tensor res, Tensor inj, " + "Tensor nw," + " float eps, Tensor wd, Tensor wu, Tensor(a!) res_out, Tensor(b!) blk_o" + "ut, Tensor(c!) inj_out, Tensor(d!) xn, Tensor(e!) sq, Tensor(f!) dpart" + ", Tensor(g!) bar, Tensor(h!) seq, int[] ar, int[] lora, int[] hb, int " + "rank, Tensor? dbg, int full, Tensor? ox, Tensor? ocodes, Tensor? ohigh" + ", Tensor? oscale, int ofmt, Tensor? gz, Tensor? gw, float geps, Tensor" + "(i!)? gscr, bool local_schedule=False) -> ()"); + ops.impl("sm70_hcx_batch_out", torch::kCUDA, &sm70_hcx_batch_out); + ops.def( + "sm70_hcx_batch_packed_out(Tensor p0, Tensor? p1, Tensor res, Tensor " + "inj, " + "Tensor nw," + " float eps, Tensor wd, Tensor wu, Tensor(a!) res_out, Tensor(b!) blk_o" + "ut, Tensor(c!) inj_out, Tensor(d!) xn, Tensor(e!) sq, Tensor(f!) dpart" + ", Tensor(g!) bar, Tensor(h!) seq, int[] ar, int[] lora, int[] hb, int " + "rank, Tensor? dbg, int full, Tensor? ox, Tensor? ocodes, Tensor? ohigh" + ", Tensor? oscale, int ofmt, Tensor? gz, Tensor? gw, float geps, Tensor" + "(i!)? gscr, bool local_schedule=False) -> ()"); + ops.impl("sm70_hcx_batch_packed_out", torch::kCUDA, + &sm70_hcx_batch_packed_out); ops.def( "sm70_dmv13_out(Tensor x, Tensor[] codes, Tensor[] high, Tensor[] scale, " "Tensor(a!)[] out, int[] fmt, int[] n, int K, int split, int warps, " @@ -396,6 +418,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int warps, Tensor(b!) ws, Tensor(c!) cnt, Tensor? sgate=None) -> ()"); ops.impl("gguf_dense_segments_sm70_out", torch::kCUDA, &gguf_dense_segments_sm70_out); + ops.def( + "gguf_dense_segments_batch_sm70_out(Tensor x, Tensor[] codes, Tensor[] " + "high, " + "Tensor[] scale, Tensor(a!)[] out, int[] fmt, int[] n, int k, int split, " + "int warps, Tensor(b!) ws, Tensor(c!) cnt, Tensor? sgate=None) -> ()"); + ops.impl("gguf_dense_segments_batch_sm70_out", torch::kCUDA, + &gguf_dense_segments_batch_sm70_out); ops.def( "gguf_shared_gate_up_sm70_out(Tensor x, Tensor[] gate, Tensor[] up, " "int[] fmts, Tensor wg, Tensor(a!) h, Tensor(b!) sg, Tensor(c!) ws, " @@ -432,6 +461,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int lanes_per_row=16, bool bank_aware=False) -> ()"); ops.impl("gguf_dp4a_gate_up_sm70_out", torch::kCUDA, &gguf_dp4a_gate_up_sm70_out); + + ops.def( + "gguf_dp4a_gate_up_grouped_sm70_out(Tensor(a!) out, Tensor activation, " + "Tensor ids, Tensor gate, Tensor up, int source_type, bool bank_aware, " + "Tensor route_tokens, Tensor order, Tensor original_ids, Tensor mode) -> " + "()"); + ops.impl("gguf_dp4a_gate_up_grouped_sm70_out", torch::kCUDA, + &gguf_dp4a_gate_up_grouped_sm70_out); + ops.def( + "gguf_dp4a_gate_up_grouped_batch_sm70_out(Tensor(a!) out, Tensor " + "activation, " + "Tensor ids, Tensor gate, Tensor up, int source_type, bool bank_aware, " + "Tensor route_tokens, Tensor order, Tensor original_ids, Tensor mode) -> " + "()"); + ops.impl("gguf_dp4a_gate_up_grouped_batch_sm70_out", torch::kCUDA, + &gguf_dp4a_gate_up_grouped_batch_sm70_out); + ops.def( + "gguf_dp4a_raw_down_grouped_sm70_out(Tensor(a!) out, " + "Tensor(b!) projected, Tensor input, Tensor ids, Tensor order, " + "Tensor probabilities, Tensor weights, int source_type, int left) -> ()"); + ops.impl("gguf_dp4a_raw_down_grouped_sm70_out", torch::kCUDA, + &gguf_dp4a_raw_down_grouped_sm70_out); + ops.def( "gguf_dp4a_lut4_gate_up_sm70_out(Tensor(a!) out, Tensor activation, " "Tensor ids, Tensor gate, Tensor gate_stats, Tensor up, Tensor up_stats, " diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index 66c5bc04ff..453b1b336d 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -38,12 +38,12 @@ identify initialization/hash code; they do not evaluate its semantics. | `vllm.config.gdn_state.GdnStateConfig` | kernel_config.gdn.state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L39); graph_options: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L52) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L23) | | `vllm.config.gdn_state.GdnStateTraceConfig` | observability_config.gdn_state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L99); compile_ignored_aliases: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L142) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L59) | | `vllm.config.kernel.IrOpPriorityConfig` | kernel_config.ir_op_priority | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L84) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L67) | -| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1044) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) | +| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1119) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L645) | | `vllm.config.kernel.Sm70AwqConfig` | kernel_config.sm70_awq | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L338); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L317) | | `vllm.config.kernel.Sm70Fp8Config` | kernel_config.sm70_fp8 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L433); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L369) | | `vllm.config.kernel.Sm70GgufConfig` | kernel_config.sm70_gguf | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L498) | | `vllm.config.kernel.Sm70NvFp4Config` | kernel_config.sm70_nvfp4 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L277); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L236) | -| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L608) | +| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L634) | | `vllm.config.kv_events.KVEventsConfig` | kv_events_config | __post_init__: [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L50) | [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L11) | | `vllm.config.kv_transfer.KVTransferConfig` | kv_transfer_config | compute_hash: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L75); __post_init__: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L93) | [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L23) | | `vllm.config.legacy_inputs.LegacyInputs` | kernel_config.sm70_awq.legacy, kernel_config.sm70_fp8.legacy, kernel_config.sm70_moe.awq.legacy, kernel_config.sm70_moe.fp8.legacy, kernel_config.sm70_moe.mxfp4.legacy, kernel_config.sm70_moe.nvfp4.legacy, kernel_config.sm70_nvfp4.legacy, scheduler_config.sm70_inputs | value: [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L39) | [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L15) | diff --git a/docs/design/flashnext_abliterated_9683_20261014.md b/docs/design/flashnext_abliterated_9683_20261014.md new file mode 100644 index 0000000000..ca5181649e --- /dev/null +++ b/docs/design/flashnext_abliterated_9683_20261014.md @@ -0,0 +1,111 @@ +# Flash-Next abliterated deployment on 9683 + +This deployment reuses the source and model contract of the +[qualified local deployment](flashnext_abliterated_local_20261014.md) on the +9683 host. It combines integration `6f3557c973`, decode `46f5c8c197` and the +qualified prefill implementation `3faaa58d08`. The copied source package is +verified against 2862 source/native file hashes. Native components are the +normal CMake build described in the local report, loaded from the deployed +package through their normal loaders; no private kernel overlay or preload is +required. The shared remote Python environment is not modified. + +## Runtime and model contract + +- Ubuntu 24.04, Python 3.12.3, Torch 2.10.0+cu128, CUDA toolkit 12.8. +- Four V100-SXM2-16GB GPUs with NVLink, physical GPUs 0–3, TP4, 46 GiB RAM. +- SC117 Flash-Next GSQ-RCO abliterated IQ3_S, revision + `eec4e10a2c29b440a2ae85591fa11569ce752963`. +- 262144 total input/output-token capacity, 8192-token chunks, one active + sequence; thinking, prefix caching and Qwen3-Coder tools enabled. +- T0.6, top-p0.95, top-k20; matching stock BF16 MTP with the existing Q8_0 + preparation route, greedy MTP4 and FULL M5 target graphs. +- Experts retain 80% GPU / 20% host placement. PLE remains on disk, target + history remains host E4M3 with encoded HOT32768, draft FP16 HOT16384. +- `QSA_HOST_INDEXER_HISTORY=true` also places compressed indexer history on + host, freeing about 230 MiB device memory. The launcher default remains + `false` for the separately measured local deployment. +- Private runtime caches, GPU ownership locks and a process-group memory guard + are retained. The guard requires 5 GiB host memory available, matching the + existing host's memory-protection threshold. + +Correctness support and fallbacks remain those of the existing implementation. +This qualification does not cover concurrent execution, other TP sizes or +quantization tiers, or every context length within the configured capacity. + +## Download and public-service transition + +The new first shard is downloaded directly on the destination from +`https://hf-mirror.com`, with real DNS and sockets bound to the physical NIC. +The host also has a transparent proxy; clearing proxy environment variables +alone is insufficient. The physical socket endpoints are retained as evidence. +The attempted machine-to-machine weight transfer is stopped at the user's +request. The final first-shard checksum is +`40f369eca301ec6ea318dbb671b3132c26a2f0b89aa5458a7051b33215ee296c`. + +The second shard is byte-identical to the existing stock shard, verified as +`316b46f3a2dbd68c900f43136ab9449f9dcc3725dfd8c794847c204bc161e113`. +It is retained through an owned hard link without another download or disk +allocation. The new first shard uses the second disk because the root disk is +nearly full. Existing model weights and source remain available for rollback. +The matching compact BF16 MTP checksum remains +`d50d273a87f51d9b7de3142a3f93f2c2823cd36a6fbf0f53f08303ce33029c17`. + +The old Strata unit owns this host's public API and GPU locks. It is stopped +only after the replacement model has downloaded and passed its checksum. The +new model uses the existing public tunnel and an authenticated proxy, with the +same credential supplied to the backend. Credentials remain outside Git. +Two duplicate tunnel clients are reduced to one, eliminating their repeated +control-channel disconnects. + +## Verification + +All shipped native components load through their normal package loaders. +Twenty-one focused GPU cases pass on the destination, covering encoded hot +storage, changed graph inputs, missing-op fallback and misaligned rejection +(22 unaffected cases deselected). The existing local qualification retains its +81-case evidence separately. + +The first configuration passes arithmetic and tools but runs out of device +memory on its first long prefill: PLE convolution needs a 144 MiB output +allocation with about 92 MiB free. The existing host-indexer setting addresses +this without reducing context capacity, expert placement, MTP depth or the +8192-token chunk. Three focused cache-budget tests pass, including the existing +host-indexer pool case. The first cold startup also exceeds the controller's +20-minute readiness timeout; its model process continues compiling normally +and completes without being restarted by that timeout. These failed attempts +and their raw logs are retained. + +With host indexer history enabled, arithmetic, automatic tool calls and replies +pass with natural EOS. Cold-prefix prefill uses the same 32000-word/seed-9683 +contract, two warmups and three measured trials, early unique nonces and no +profiler. Actual measured inputs are 34688, 34606 and 34631 tokens. Median TTFT +is **6.944 s / 4987.2 input tokens/s**. After applying the existing host's CPU +isolation pattern only to owned API/worker descendants, a matching repeat gives +**6.945 s / 4982.5 tokens/s**: no measurable prefill gain is claimed. These +measurements are close to, but do not establish the strict 5000-tokens/s target. + +The third identical 32774-input prefix request proves 31824 cached tokens; +total request time falls from 11.476/10.886 s to 4.621 s. Retrieval/arithmetic +tasks at 32761, 131070 and 260977 input tokens pass answer, EOS and repetition +checks (196, 199 and 210 output tokens). The largest completes with 261187 +total tokens in 310.776 s, without OOM or a memory-guard intervention under +the final configuration. These are task/capacity checks, not cold throughput +baselines at each length. + +Public authentication rejects an unauthenticated model query; authenticated +discovery and streamed arithmetic succeed with the correct 1517 and EOS. +The initial external generation queues behind the long-context check and +takes 49.689 s to its first chunk; after the owned test queue is empty, the +same short check reaches its first streamed chunk in **1.050 s**, completing +in 1.341 s. First HTTP/SSE chunks, first generated tokens and first visible +answer text remain separate concepts. Single active-sequence capacity means +simultaneous Agent requests can still queue. + +The new model and proxy units are enabled for startup; the old Strata unit is +disabled with its definition and prior enabled state retained for rollback. +The CPU isolation helper is retained in the deployment wrapper for subsequent +starts. Production development cache-reset endpoints remain disabled. + +The external `flashnext-abliterated-9683-20261014` report retains checksums, +source and native identity, failed/abandoned attempts, GPU tests, runtime logs +and request responses. Use its handoff for the live service paths and endpoint. diff --git a/docs/design/flashnext_abliterated_local_20261014.md b/docs/design/flashnext_abliterated_local_20261014.md new file mode 100644 index 0000000000..f77658466a --- /dev/null +++ b/docs/design/flashnext_abliterated_local_20261014.md @@ -0,0 +1,123 @@ +# Local Flash-Next IQ3_S abliterated deployment + +This deployment combines the latest integration base `6f3557c973`, the decode +implementation from PR #1204 at `46f5c8c197`, and the qualified prefill +implementation from PR #1205 at `3faaa58d08`. Later prefill experiments remain +outside this deployment. + +The model is +[SC117/Qwen3.8-Flash-Next-GSQ-RCO-abliterated-GGUF](https://huggingface.co/SC117/Qwen3.8-Flash-Next-GSQ-RCO-abliterated-GGUF), +IQ3_S, revision `eec4e10a2c29b440a2ae85591fa11569ce752963`. Both GGUF shards +are downloaded through `https://hf-mirror.com`, with resumable ranges and the +repository's SHA256 values checked before serving. The initial environment-only +proxy bypass did not bypass the host's transparent TUN. Subsequent requests use +real DNS and sockets bound to the physical NIC; the routing mistake and corrected +connection evidence are retained in the external report. +The matching stock BF16 MTP tensors are copied from the existing local +checkpoint unchanged and verified against the previously qualified compact +checkpoint hash. The runtime prepares its existing Q8_0 MTP expert route. +The repository's standalone Q8_0 GGUF draft file is also retained locally. + +## Runtime contract + +Use [the launcher](../../examples/serve_flashnext_gguf_256k.sh) from its owning +source checkout and an environment with the matching native extensions. + +- Four V100-SXM2-16GB GPUs, physical GPUs 4–7, TP4, one active sequence. +- Python 3.12.13, Torch 2.10.0+cu128, Triton 3.6.0, CUDA toolkit 12.8. +- Capacity is **262144 total input and output tokens**, with chunk size 8192. +- Prefix caching, Qwen3 reasoning and Qwen3-Coder tool calling enabled. +- Temperature 0.6, top-p 0.95, top-k 20; thinking enabled. +- MTP4, greedy draft, FULL target graphs at M5, existing draft single graph. +- Original experts retain 80% GPU / 20% host placement. PLE remains on disk. +- Authoritative target KV uses host E4M3. Encoded target hot capacity is 32768; + the independent draft cache retains FP16 and 16384 tokens. +- The accepted grouped-prefill and dense restoration policy, cold-expert DMA, + balanced chunks and host-KV zeroing skip are enabled through existing paths. +- Rejected HC/PLE, indexer, tile and larger-chunk screens are not added. +- Native components are built by normal CMake targets into the source package. + No private native sidecar, source overlay or preload is required. +- Compiler caches and the service's memory guard are owned by this deployment. + +These settings preserve existing runtime support and fallback boundaries. +They do not qualify a different TP size, concurrent requests, all quantization +tiers, or every length within the configured capacity. + +## Reproduce + +Set an API key and the local paths before launching: + +```bash +export MODEL_FILE=/path/to/IQ3_S/Qwen3.8-Flash-Next-GSQ-RCO-abliterated-IQ3_S-00001-of-00002.gguf +export MTP_DIR=/path/to/mtp-stock-bf16 +export CACHE_DIR=/path/to/dedicated-runtime-cache +export PYTHON_BIN=/path/to/venv/bin/python +export CUDA_HOME=/path/to/cuda-12.8 +export VLLM_API_KEY='your-key' +export GPU_IDS=4,5,6,7 CPU_SET=24-47,72-95 +export QSA_HOST_INDEXER_HISTORY=false +bash examples/serve_flashnext_gguf_256k.sh +``` + +Acquire the selected GPUs' locks and start the memory guard as part of the +service manager. The local deployment retains the existing user service and +public tunnel, with a dedicated internal API port. The existing bridge now +forwards to the verified model instance. Read the tunnel client's published +endpoint when connecting: its assigned public port can differ from the port +requested in the preset configuration. Credentials stay outside the source +repository. + +`QSA_HOST_INDEXER_HISTORY` accepts `true` or `false` and defaults to `false`, +preserving this local measurement's configuration. `true` selects the existing +host-backed compressed-indexer cache, reserving more device memory for prefill +temporaries. The [9683 deployment](flashnext_abliterated_9683_20261014.md) +records the separate qualification of that setting. + +## Evidence + +The prefill dependency's original-model TP4 screen measured **5010 tokens/s** +at 34688 input tokens, median TTFT 6.907 s. The decode dependency's original +32K screen measured conservative mean verifier cost 15.494 ms. These are +dependency results on the 9683 host, not performance measurements of this +abliterated model on the local machine. + +The combined build passes 81 focused GPU cases covering encoded/protected KV, +expert-cache swaps, grouped-prefill boundaries and segment-prefill restoration. +All twelve normal native components are built and their library search paths +are audited. All model files pass their complete SHA256 checks. Worker logs +confirm native encoded QSA, Q8_0 MTP expert storage and the existing GDN/PLE paths. + +Five unprofiled cold-prefix requests use two warmups and three measured trials, +with the same 32000-word/seed-9683 workload and an early per-request nonce. +Production development cache-reset routes remain disabled; the unique first +block prevents prefix reuse. Measured inputs are 34688, 34606 and 34631 tokens. +Median first-choice latency is **6.805 s**, or **5088.8 input tokens/s**. Server +prefill time is recorded separately. This is the local abliterated model's +measurement, rather than a projected gain from the dependency results. + +Eight prefix-warm, unprofiled 32774-input / 192-output requests at production +sampling measure **6.454 ms/output token**, or **154.9 output tokens/s**. +Mean decode time is 1.233 s, excluding 0.893-s TTFT. The whole-round HTTP +estimate is 18.128 ms with 2.838 emitted tokens per round. It includes draft and +scheduling and must not be labelled an isolated GPU verifier timer. The +recorded affinity changes affect only this service's API/worker descendants. + +Arithmetic and automatic tool-call/follow-up checks pass with natural EOS. +Retrieval/arithmetic tasks at 32761, 131070 and 260977 input tokens pass all +answer, EOS and repetition checks, with 181, 165 and 188 output tokens. The +largest case uses 261165 total tokens and finishes in 286.959 s without OOM +or a memory-guard intervention. These are capacity and task checks, not cold +throughput baselines at each length. + +Prefix-cache counters prove 31824 cached tokens on the third identical +32774-input request after the long-context checks. Total request time falls +from 8.921/9.304 s to 2.731 s. A separate 18282-token repeated prompt does not +produce a hit in two trials; that observation is retained. Prefix caching uses +the existing Mamba alignment/admission policy and finite cache capacity. This +evidence establishes reuse for the measured workload, not every prompt length +or guaranteed residency after other requests. + +Public authentication, model discovery and streaming arithmetic are checked +through the existing tunnel before delivery. Raw download journals, checksums, +builds, counters and model responses are retained in the task's external +`flashnext-abliterated-local-20261014` report directory. diff --git a/docs/design/flashnext_decode_verifier_16gb_20261014.md b/docs/design/flashnext_decode_verifier_16gb_20261014.md new file mode 100644 index 0000000000..684c98eaee --- /dev/null +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -0,0 +1,407 @@ +# Flash-Next verifier diagnosis on four 16GB V100 GPUs + +The objective is target verification below 16 ms at approximately 32K input +tokens, while retaining emitted-token speed and output quality. Verification +includes target forward, target sampling and state update. Draft proposal is +reported separately. The final opt-in encoded target hot cache measures an +all-round mean of 15.482–15.494 ms at 32K, including conservative bounds for +two missing graph starts. This meets the mean objective; p90 and p99 remain +17.340 and 31.689 ms. Ordinary decode and near-256K completion also pass the +measured cases described below. + +## Fixed workload + +The reference source is `6f3557c973b4602905037e086fd4acba9b0d867f`. +Hardware is four V100-SXM2-16GB GPUs with a full NV2 mesh, driver 570.211.01, +CUDA 12.8, Torch 2.10.0+cu128 and Python 3.12.3. The model is Flash-Next +GSQ-RCO IQ3_S with Q8_0 MTP4, greedy draft sampling, TP4 and one request. +Target sampling is temperature 0.6, top-p 0.95, top-k 20, seed 20261014; +thinking and EOS remain enabled. + +Context capacity remains 262144 **total input and output tokens**, chunk size +8192, prefix caching enabled, and FULL target decode graphs at M5. Expert storage retains +the configured 80% GPU / 20% host placement. Target QSA history uses host +E4M3; the reference has an 8192-token device hot cache and the configuration +candidate has 16384 tokens. Draft history uses host FP16. Device hot values +remain FP16 for both history types; the larger cache adds about 104 MiB per +rank including draft storage. +PLE remains on disk with its native lane and 512-MiB CPU row cache. +Measured prompts contain 574 and 32774 tokens, with 192 output tokens. +These measurements do not qualify 256K speed or concurrent requests. + +## Reference measurements + +Ordinary streaming HTTP excludes TTFT from decode. Metrics count actual +draft-verification steps, so the round estimator divides by draft count, +without subtracting a prefill step. Retain raw observations and first-run +variation; the HTTP estimate is not an independent verifier timer. + +| Input tokens | Mean ms/output token | Estimated ms/round | Tokens/round | +| --- | ---: | ---: | ---: | +| 574 | 9.110 | 18.710 | 2.054 | +| 32774 | 8.721 | 21.256 | 2.450 | + +CUDA-only Nsight graph-node capture observes 1120 target kernels on every +rank. Short/long target counts are 94/76, versus independent committed +verification counts of 93/75; retain the extra pipeline-end replay. Draft +graphs contain 70/71 kernels and 207/210 kernels. Requests are separated by +their prefill boundary, and TP replays are aligned by ordinal. + +The following means close the same GPU0 target-start to next-target-start +window at 32K. Node tracing perturbs latency; this is composition evidence. + +| Phase | Mean ms | +| --- | ---: | +| Target graph activity envelope | 22.124 | +| Target end to first draft graph | 0.862 | +| Dependent draft graph chain | 3.895 | +| Draft end to next target graph | 1.173 | +| Complete round | 28.053 | + +Expert copies contribute 2.437 ms of target kernel service, expert publish / +CPU-plan wait 1.206 ms, and QSA host gather 1.990 ms. Dense MMA contributes +2.272 ms and 94 complete HC boundaries contribute 2.151 ms. Service can +overlap; these values must not be added into an endpoint wall table. + +Explicit synchronized step timing retains every TP worker. At 32K, the two +reference requests have mean per-round maximum verifier CPU walls of +20.354 / 17.701 ms, p90 25.496 / 20.319 ms, and p99 40.704 / 25.687 ms. +Target sampling is approximately 0.42 ms and state update 0.02 ms. +Event collection, fences and interval-one logging perturb endpoint latency; +ordinary HTTP remains the accepted emitted-token speed measurement. + +## Candidate and operator evidence + +The existing expert copy kernel changes from 16 operation programs with +8 word splits, BLOCK1024/W8, to 4 operation programs with 32 word splits +and 3 parallel projections, BLOCK256/W4. This increases parallel work for +small miss plans without serializing all three projections in each CTA. Whole-word +storage, swap ordering and masked tails retain their exact byte contract. + +The paired TP4 screen uses real per-expert byte geometry from layers 0, 1, +17 and 47, with production 410 device / 136 host slot counts. All six +policies pass eager and changed-input graph equality over complete host and +device banks, including guards. The initial 4-by-64 word split improves +common one/two-miss cases roughly 6–20%, with a 6.4% regression in one +eight-miss case. The projection-parallel screen additionally improves paired +miss plans by 0.5–17.5% versus that first candidate, while some swap cases +regress. End-to-end admission is required. Word lengths 255/256/257 and +operation counts around 4 and 16, through 128, also pass four-rank byte checks. + +For QSA direct resolution (STAGE0), values that cannot be installed are +read later by the authoritative native attention reader. The candidate +omits the unused host read/decode and mixed-tile cached-value read, and +skips the entire value-materialization block when no page is installed. It +preserves tags, ownership, counters, resolved indices and consumed values. +Staging modes 1/2 retain their original value-read masks; staging writes +complete before the same release of cache metadata. + +Correctness support retains the existing valid page4/D256 geometry and +FP16/E4M3 storage. The changed subset is direct resolution with failed +installations or mixed hit/miss tiles. Shape admission and existing native +and staged fallbacks are unchanged. No numerical tolerance is relaxed. +The first host KV suite passes 104 cases on both reference and candidate, +including protected-capacity fallback and changed-input graph replay. The +second candidate passes 108 host KV cases and 7 expert cache cases, including +probe-pointer signed/unsigned wrap, graph replay and decode/prefill interleaving. +The step-profiler suite passes 13 cases on GPU0. + +The first candidate's 32K second request improves maximum verifier CPU-wall +mean from 17.701 to 16.704 ms, p90 from 20.319 to 18.388 ms, and p99 from +25.687 to 19.113 ms. Its first request is 20.598 ms with an 87.139-ms maximum, +versus the reference 20.354-ms mean; retain this first-request variation. +The second candidate's first/second 32K means are 20.054 / 16.864 ms; its +second-request p90/p99 are 18.373 / 25.294 ms. This does not establish a +further gain or satisfy 16 ms. Ordinary candidate HTTP qualification remains +pending while optimization continues. + +Existing tests had compared the full 32-row workspace against a smaller +active request and passed capacity-sized metadata to attention. The +original kernel reproduces the failure. Tests now slice the active rows; +the original and candidate pass the same corrected oracle. + +## Reproduce and retained negative paths + +Run the installed source benchmark on an idle, locked TP4 GPU set: + +```bash +.venv/bin/python -m torch.distributed.run --standalone --nproc_per_node 4 \ + benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py --output copy.json +.venv/bin/python -m torch.distributed.run --standalone --nproc_per_node 4 \ + benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py \ + --boundary-only --samples 1 --repeats 2 --output boundary.json +``` + +Use the prepared environment's Python and its normal native artifacts. +Raw Nsight report, SQLite, HTTP responses, per-rank logs and copy JSON are +retained in the task's external `flashnext-decode-verifier-9683-20261014` +report directory, outside Git. + +Optional NVTX scopes failed because the Python package is absent; CUDA-only +node tracing succeeds. A concurrent prefill compiler exhausted host headroom +and memguard terminated an early capture; the user chose compiler priority. +The serving CLI does not expose `--observability-config`; the existing +captured MTP profiling compatibility flags enable the diagnostic instead. + +Initial direct-resolve screens retained one KV epoch across replays. They +measure a protected-cache regime: selected pages remain pinned and later +misses fall back instead of installing. They do not represent normal decode +churn. The corrected fixture advances the epoch before every resolve, +matching the protection transition of a decode KV write. + +With advancing epochs, the adversarial FP8 M5 churn screen measures W4/W8 +at 693 / 660 us, while shared/overflow selections regress about 10% at W8. +Eight-token tiles improve that churn case only about 2%; four-token tiles +regress both shared selections and churn. Retain the existing W4, 16-token +tile and per-page CLOCK atomics. Register counts are 127 for W4/T16, 76 for +W8/T16 and 72 for W4/T8, with no spills. This is a synthetic operator screen; +register reduction alone does not qualify model speed. No new tiling or warp +startup switches are added. + +An exact-word HC gather experiment encoded Half storage words as integers in +FP32 and used the existing peer reduction. Every Half bit pattern and changed +graph passed exact equality, but the complete M5 boundary regressed from +69.824 to 79.296 us; the model's second 32K verifier mean was 17.472 ms. The +helper was reverted. Collective substitution alone is not a speedup. + +The existing 32768-token hot-cache setting also failed admission: the first +real request needed a 144-MiB allocation with only 75.94 MiB free. It retained +the expert ratio and context capacity but cannot run on this 16GB deployment. +The 16384-token setting runs successfully and measures first/second 32K +verifier means of 18.435 / 16.059 ms, with second-request p90/p99 of +17.209 / 18.848 ms. It is a separate configuration screen and still misses +16 ms. Arithmetic, Python evaluation, two 32K retrieval tasks and a 1809-token +CPU explanation pass their task/EOS gates; the explanation has no six-word +phrase repeated more than twice. Ordinary throughput qualification is pending. + +The two target-graph NCCL gathers occur after the early PLE boundary. PLE +materializes the pending MoE residual before adding its result, then calls +standalone HC `mix()`. Retaining the six-dimensional HCX fusion pack at that +boundary forces unpacking and two gathers, because standalone LL projection +consumes the existing five-dimensional pack. A candidate leaves partial +residual ownership intact and prepares the LL bank at PLE boundaries instead +of the unused HCX bank. This changes one target boundary on the measured +model and allocates no duplicate bank. HCX schedules at other boundaries, +partial residual ownership, normalization and request state were unchanged. + +The actual layer-1 checkpoint screen on four ranks measures complete M5 +boundaries at 76.672 / 19.140 us for reference / LL storage. M1/M8/M20 are +also faster; M21 and M256 retain fallback with exact outputs. Native arithmetic +uses the existing FP16 HC error gate, rtol 0.003 and atol 0.003, fixed before +the screen. Both variants replay changed inputs exactly against their own +eager results, with finite outputs. CPU dispatch/lifetime tests pass 21 cases; +the corrected fallback suite passes 17 GPU cases across two rank orders, +both storage layouts, M5/M20/M21 and M512. The model screen retains HOT8192 +to separate this source change from cache tuning. + +That model screen rejects the PLE storage change. First/second 32K verifier +means are 18.077 / 17.048 ms; the second-request p90/p99 are 19.262 / 23.022 ms. +Committed draft steps rise to 91 for 192 output tokens, versus 75 in the +preceding HOT8192 candidate. Four integer tasks pass, but the CPU explanation +remains in reasoning at the fixed 2048-token limit. This does not establish a +quality regression by itself, but gives no qualification for promotion. +The runtime and dispatch-test changes are reverted to the preceding candidate. +Retain the corrected fallback oracle and its passed GPU coverage. The large +operator speedup and node-trace NCCL waits did not predict lower live verifier +cost; do not promote this route based on the microbenchmark alone. + +The old fallback oracle omitted the physical rank and simulated peers by +scaling an already mixed local result. That does not model each peer's +different input columns and nonlinear gates; it fails the unchanged HCX +fallback too. The corrected oracle computes remote projections and columns, +includes the actual local packet, and retains exact assertions and changed +input graph checks. + +An asynchronous final verifier replay can finish after the HTTP stream. +The original before-counter read occasionally assigned that replay to the +next request. Retained API POST boundaries align all four workers and keep +every final replay with its preceding request; aligned counts differ from +independent committed draft metrics by at most one. Raw and aligned ledgers +are both preserved. New event probes wait for all-rank counters to remain +stable for 200 ms between requests. This does not change measured decode +timestamps or ordinary production requests. + +## Ordinary whole-graph measurements and completeness + +CUDA graph-level tracing records graph envelopes without per-round CUDA +Events or fences. The verifier window runs from the earliest target graph +GPU start to the latest `_postprocess_mamba_align_kernel` end across TP4. +Source inspection establishes `_post_update_kernel`, then +`_scatter_num_accepted_kernel`, snapshot copy and Mamba alignment as target +sampling/state commit. Draft preparation follows this boundary. CPU graph +launch-to-draft time includes previously queued work and is not an isolated +verifier measurement. + +The matched ordinary HTTP pair retains identical prompts, seeds, sampling, +MTP, native artifacts and warmup sequence. Reference HOT8192 versus source +candidate HOT16384 measures 574-token mean TPOT at 8.957 / 7.714 ms, and +32774-token mean TPOT at 8.331 / 8.519 ms. The latter is about 2.3% slower; +there is no unconditional emitted-token speed claim. Tokens per committed +round differ, so output TPOT and verifier latency must remain separate. + +Commit-aligned four-request traces measure first/second 32K verifier means +of 16.967 / 16.169 ms for the reference and 15.816 / 17.788 ms for the +candidate. The candidate's 86.557-ms tail remains included. Every committed +round is retained; one additional asynchronous pipeline-end graph is +recorded separately per request. + +An eight-request steady HOT16384 capture has 584 complete state commits on +every rank: each request contains one prefill commit, 71 committed verifier +rounds and one pipeline-end replay. Five graph activity records are missing +across the ranks, affecting four of the 568 committed rounds. Direct target +graph ordinal alignment is invalid for this report and its old global +statistics are withdrawn. The three commit kernels, target graphs and +draft graphs all use the same CUDA context and stream on each rank. +Associating each graph with its following commit aligns complete request +and round ledgers. For a missing start, the preceding same-stream commit +end bounds the target start from below; this conservatively includes draft +and preparation time instead of imputing or dropping the missing window. +All remaining graph records map uniquely to a commit and all independent +committed-step counts reconcile. + +The resulting eight-request all-round verifier mean lies between 16.014 +and 16.043 ms; upper-bound p50/p90/p99/max are 15.265 / 18.984 / 24.134 / +31.897 ms. This remains above the 16-ms objective. Nsight's incomplete-event +warning and four conservatively bounded rounds are retained. A fresh single +capture range is required for a fully exact distribution. Interactive rearm +returned `StopRequest` errors and did not reliably honor a new output/range; +do not mix its report with the later HTTP requests. + +HOT20480 passes the five screened task/EOS cases but measures synchronized +32K verifier means of 17.864 / 16.590 ms, slower than HOT16384 and with less +memory headroom. W2/T16 at HOT16384 uses 202 registers without spills but +regresses shared, overflow and advancing-epoch churn selections; retain W4. +No corresponding runtime or default tuning change is promoted. + +The ordinary pair passes arithmetic, Python evaluation and both 32K needle +retrieval tasks. Its open-ended CPU explanation reaches the 2048-token cap +in both arms; a candidate 4096-token follow-up also reaches its cap. These +runs do not qualify EOS or long-output drift. The earlier HOT16384 screen's +1809-token natural EOS is retained separately. A 261959-input-token capacity +request reaches the 300-second client timeout without an OOM; it does not +establish successful near-256K completion. Keep both gates explicit before +promoting the deployment configuration. + +An empty-plan expert-ID remap screen tests counts 0/1/5/10/50/128 and plan +sizes 0/1/2/4/20/256 with W4 and W1, with and without an empty-metadata branch. +All 36 cases per policy retain exact table/output bytes in eager execution +and changed-input graphs. The measured M5 empty-plan kernel changes from +1.832 to 1.992 us with the W4 branch, so no runtime change is retained. + +The 564 exactly recorded steady rounds show mean maximum individual-rank +verifier duration of 15.656 ms, versus 16.017 ms for the complete TP4 window. +GPU target-start skew has mean/p90/p99 of 1.081 / 3.543 / 6.824 ms. CPU target +enqueue skew averages 8.579 ms; GPU queue-ahead averages 11.246 ms and includes +prior work. Enqueue skew correlates with GPU-start skew (0.607), motivating +an owned-process CPU affinity screen rather than treating queued time as +additional verifier computation. + +## Encoded hot-cache model screen + +The opt-in target setting `qsa_host_kv_hot_dtype="fp8_e4m3"` retains the +encoded bytes and original per-vector FP32 scales in device hot storage. +Both the existing native QSA family and its protected Triton reader decode +to the same FP16 operands used by the decoded cache. The authoritative host +history, page ownership, request masks, tentative writes, release ordering +and staging format remain unchanged. Default hot storage remains FP16. +Encoded hot storage requires E4M3 history; incompatible history is rejected. +An unavailable `run_hot_fp8` native artifact uses the qualified protected +reader, and unaligned native hot storage is explicitly rejected. + +The optional `qsa_host_kv_draft_hot_tokens` independently bounds the draft +cache; its default None preserves the existing shared capacity policy. +The measured prototype doubles target hot capacity from 16384 to 32768 +while retaining the 16384-token FP16 draft cache. Target storage changes +from 1024 to 520 bytes per cached token. Across twelve target QSA layers, +this configuration adds about 3 MiB per rank including scales, while leaving +the expert ratio, host histories, disk PLE and total context capacity intact. + +GPU qualification passes 151 cases: the original 130 host/native tests, +18 new native/protected-reader cases over M1/M5/M20/M21/M32, selection +widths 1/15/16/17/65/2051/4096, hot capacities 48/128/32768, request masks, +causal bounds and changed-input graphs, plus unavailable-native fallback +and two alignment-rejection cases. Live hot bytes and scales match the +host source exactly. Native attention uses its existing FP16 comparison +gate, rtol 0.003 / atol 0.0002, unchanged before the candidate runs. +The first prototype's eight fallback failures are retained: its generic +reader still interpreted encoded bytes as FP16, and three fixtures used +invalid staging capacities. The corrected reader and fixtures pass. + +The native module is built with the configured target's exact NVCC/link +arguments, writing only task-owned outputs and replacing the deployed +source module atomically. Shared venv packages are untouched. Original +native SHA256 is `ed1d3a854e673cec3c646076bbc7d3b3010d5840fa4c17867d1fc445a2e41eac`; +prototype SHA256 is `76fa88cc1fe323f019c0ad6942756dde84e5396bf46ae07b00cd27dc99115fd2`. +Rebuilding the final formatted source produces SHA256 +`bf0dd515b90a2132f95ae408ab4e23bec4e0e19b81b1199f2d74b9d644b7a54b`. +All 24 affected encoded-cache, unavailable-native, alignment-rejection and +legacy E4M3-code tests pass again against that final artifact. +Rebuilding the final formatted source gives SHA256 +`bf0dd515b90a2132f95ae408ab4e23bec4e0e19b81b1199f2d74b9d644b7a54b`. +Its 24 targeted GPU regressions pass, covering encoded/protected-reader +graphs, missing-native fallback, alignment rejection and all E4M3 codes; +the preceding 151-case suite remains the broader qualification evidence. + +The TP4 M5 screen includes resolve plus native attention at 32774 input +tokens, width 2051, advancing KV epochs and alternating paired measurements. +FP16 HOT16384 / encoded HOT32768 shared-selection medians are +51.456/52.992, 43.264/45.808, 48.896/53.760 and 47.104/49.920 us across +ranks. Sixteen changing selections measure 719.776/595.888, +731.136/592.207, 713.152/598.064 and 716.224/590.608 us. Churn improves +17–20%, with shared selections 3–10% slower. This passes the predeclared +operator screen of at least 5% churn improvement and at most 10% shared +regression. Model admission uses the separate measurements below; the +repository's default hot storage remains FP16. + +The encoded configuration's eight-request model capture retains all 640 +committed verification rounds. State commits are complete on TP4; 638 +windows are exact and two missing graph starts use conservative bounds. +The all-round mean lies between 15.482 and 15.494 ms, below the 16-ms mean +objective even at the conservative endpoint. Upper-bound p50/p90/p99/max +are 14.687 / 17.340 / 31.689 / 35.538 ms. Request means are 15.209, 14.780, +15.894, 16.804, 15.303, 15.072, 16.119 and 14.772 ms: this is an average +objective, not an every-round or p99 guarantee. The preceding FP16 HOT16384 +capture under the same affinity policy has mean bounds of 16.079–16.091 ms; +the runs have different warmup histories, so this is consecutive-screen +comparison rather than an unconditional throughput gain. + +Arithmetic, Python evaluation and both 32K retrieval tasks pass with natural +EOS. The paired concise 32K CPU task passes both seeds 20261014/20261015: +the encoded run ends naturally at 1088/682 output tokens versus 1333/1348 in +the preceding FP16-hot run. Both final answers correctly explain renaming, +RAW readiness, the ROB's ordered retirement and precise exceptions, with +concrete dependency examples. Reasoning and final output have no six-word +phrase repeated more than twice. The open-ended 2048-token CPU task still +hits its cap and remains an inconclusive long-answer stress observation. + +Eight ordinary streaming requests, with capture inactive and the same 32774 +input / 192 output tokens, measure mean TPOT 7.109 ms, or 140.7 output +tokens/s. Mean decode time is 1.358 s, excluding 1.309-s prefix-warm TTFT. +The HTTP round estimator is 18.859 ms with 2.681 emitted tokens per round; +it includes draft and scheduling and is distinct from GPU verification. +All requests emit 192 tokens. This is a candidate observation, not a paired +throughput gain over the earlier runs, whose acceptance counts differ. + +The same capacity probe, with its timeout raised from 300 to 900 seconds +and unchanged 128-output limit, completes with 261959 input and 71 output +tokens (262030 total), natural EOS and the correct `SAFFRON45793` marker. +TTFT is 320.034 s and total time is 320.435 s. This establishes completion +near the 262144-total-token boundary for this request, not 256K speed or +concurrency. No OOM or memguard termination occurs. GPU memory reaches +16143 MiB per rank afterward, leaving roughly 17 MiB free: capacity passes +with very little device headroom. The earlier 300-second timeout remains +retained as an incomplete observation rather than a successful capacity run. + +The final ordinary server runs without Nsight injection or profiler CLI +configuration, with target encoded HOT32768 and draft FP16 HOT16384. After +three warmup requests, task-owned worker main and helper threads are placed +on separate physical CPU cores using the recorded affinity policy. One +32774-input / 192-output smoke measures TPOT 7.836 ms, or 127.6 tokens/s, +and estimated whole-round cost 18.252 ms. Its acceptance is 0.335, versus +the earlier eight-request screen's different acceptance history. This is +a single deployment smoke, not a new mean-verifier distribution. +The key-gated public proxy returns 401 without authentication and 200 for +the authenticated model list and generation. Arithmetic generation returns +1517 with natural EOS. Memory protection remains active. The standard +deployment launcher retains its GPU locks and now records the selected +target/draft cache settings; its original script is backed up outside Git. diff --git a/docs/design/flashnext_mtp_batch_16gb_20261015.md b/docs/design/flashnext_mtp_batch_16gb_20261015.md new file mode 100644 index 0000000000..7a52844f34 --- /dev/null +++ b/docs/design/flashnext_mtp_batch_16gb_20261015.md @@ -0,0 +1,1085 @@ +# Flash-Next MTP4 batch scaling on four 16GB V100 GPUs + +The optimization target is the MTP4 verifier, with the qualified C1 cost of +approximately 16 ms as the reference. One verification round processes five +target rows per request. Its cost includes target forward, target sampling and +state commit; draft generation is separate. HTTP time per emitted token and +prefill throughput cannot substitute for this measurement. + +## Workload and evidence boundaries + +Integration base: `onecat/main`, `6f3557c973b4602905037e086fd4acba9b0d867f`. +The owned branch starts from the qualified deployment composition +`699c864d3c3a841b8626b657a1c8f236457ab1b5` (decode #1204, prefill #1205 and +deployment #1206). Native libraries were rebuilt from that source in an owned +CMake directory. All 12 packaged libraries have recorded SHA256 values and +no dependencies under private build directories or another worktree. + +Hardware: local GPUs 4–7, Tesla V100-SXM2-16GB, TP4, driver 580.178.04, +CUDA 12.8, Torch 2.10.0+cu128, Python 3.12.13. The model is SC117's +Qwen3.8-Flash-Next-GSQ-RCO-abliterated IQ3_S with the matching stock MTP +weights converted to Q8_0. Context capacity is 262144 total tokens, prefix +caching, aligned recurrent state, disk PLE and host QSA E4M3 history. +Experts retain the required 80% device / 20% host placement. +Sampling remains temperature 0.6, top-p 0.95, top-k 20, seeds starting at +20261014; thinking and natural EOS remain enabled. Draft length is four with +greedy draft sampling. The measured long prompt has 32774 tokens. + +The earlier 16 ms evidence used the original GSQ-RCO model on 9683 and kept +compressed indexer history on device. The first local batch diagnostic uses +host compressed indexer history and a 512-token prefill chunk to fit the +larger state pool. Its absolute C1 cost must not be substituted for the +qualified 16 ms result. The original-layout C1 control is measured separately. + +The owned-source original-layout C1 node capture confirms this distinction: +36 interior complete rounds have a 16.009 ms median and 16.234 ms mean, with a +16.243 ms TP4 envelope mean. Trace service on the critical rank is 3.295 ms +dense/shared projections, 2.882 ms expert projections, 2.322 ms HC/fused +communication and 1.985 ms QSA history/attention. The inactive-capture HTTP +control emits 143.1 tokens/s in its common decode interval with acceptance +length 2.838. Nsight injection is still present in that control; it is not a +fresh uninstrumented production result. Full per-round data and exclusions +remain in `c1-original-layout-node-breakdown.json`. + +## Initial verifier diagnostic + +Synchronized worker events retain all four TP ranks. Align the same step +ordinal, require `tokens=5*C` and `drafts=4*C`, and select the rank with the +largest verifier event sum. Stage values below come from that same rank. +Event profiling adds synchronization and is diagnostic evidence, not an +unprofiled production speed result. + +| Actual concurrent requests | Complete rounds | Verifier mean ms | Target forward ms | Sampling ms | State commit ms | Draft ms | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 1, repetition 0 | 112 | 27.251 | 26.829 | 0.405 | 0.017 | 4.707 | +| 1, repetition 1 | 112 | 27.247 | 26.824 | 0.405 | 0.017 | 4.706 | +| 2, repetition 0 | 93 | 69.715 | 69.262 | 0.433 | 0.020 | 7.230 | +| 2, repetition 1 | 97 | 69.580 | 69.127 | 0.433 | 0.019 | 7.232 | +| 4, first actual overlap | 32 | 102.197 | 101.650 | 0.524 | 0.024 | 10.583 | + +The C4 row is from the first warmup batch's actual four-request overlap. Its +later cached batch never reached M20, and the HTTP common-window check failed. +That failed attempt is retained and cannot be reported as sustained C4. +C2 verifier inflation is about 2.56x; C4 is about 3.75x. Nearly all of it is in +target forward. GPU kernel attribution and paired candidate timings are needed +before selecting a performance default. + +## Full-width MTP4 batch measurements + +The compact-state/tiled-kernel candidate with grouped shared-key scoring +produced sustained full-width verifier rounds at all five concurrency levels. +The existing scheduler option `--mixed-prefill-min-tokens 512` was fixed for +this pass: its adaptive floor of 16 previously starved new prefills, so C6/C8 +never overlapped. This is a lab configuration change, not a new scheduler +default. Only rounds with five target rows and four draft rows per request +are included below. All four ranks must be present. + +| Concurrency | Complete rounds | Verifier mean ms | p50 ms | p90 ms | p99 ms | Forward ms | Sampling ms | State ms | Draft ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 27.422 | 27.268 | 28.019 | 29.934 | 26.997 | 0.408 | 0.017 | 4.798 | +| 2 | 191 | 52.964 | 52.745 | 53.764 | 56.661 | 52.510 | 0.435 | 0.019 | 6.760 | +| 4 | 206 | 93.346 | 92.797 | 95.453 | 101.870 | 92.796 | 0.527 | 0.023 | 10.753 | +| 6 | 182 | 138.721 | 138.250 | 140.552 | 150.466 | 138.115 | 0.580 | 0.026 | 15.069 | +| 8 | 168 | 185.889 | 185.226 | 190.645 | 196.910 | 184.995 | 0.864 | 0.030 | 19.122 | + +With grouped scoring disabled, C2 was 60.521 ms and C4 was 102.246 ms under +the earlier tiling pass. The first control C2 was approximately 69.6 ms. +Those C4 controls have fewer overlap rounds and are not a paired sustained-C4 +production result. Per-request verifier share above is 27.4/26.5/23.3/23.1/23.2 +ms. It remains above the qualified 16 ms device-indexer C1 reference. + +Node attribution localizes the remaining cost to compressed-indexer scoring, +rather than main selected-history attention. In the short retained trace, +C1/C2 indexer GPU service is 13.644/28.931 ms on the critical verifier rank; +the native shared-key reader accounts for 13.008/28.292 ms. Main selected +history attention is only 0.663/1.033 ms. These are traced service sums from +64 C1 and seven C2 rounds, not the event-profile absolute timings above. +The trace retains a separate wall residual and incomplete-round exclusions. + +## Candidate changes and retained fallbacks + +* QSA prefill staging now checks per-request query length, rather than treating + every batch above 32 total rows as prefill. MTP4 C8 is M40 with query length + five. The existing native reader handles its 32+8 fragments. True prefill + retains staging; unsupported direct inputs retain the protected reader. +* `sm70_hcx_batch_max_m` defaults to 8. An explicit limit up to 40 uses existing + native eight-row HCX tiles. The limit belongs to each module, not mutable + shared runtime policy. Larger batches keep projection/reduction fallbacks. +* `sm70_gguf.small_m_dp4a_max_m` defaults to 20. An explicit limit up to 40 + uses existing native experts in fragments of at most 20 rows. Fragment size + also respects the cache's resident slots and miss ring. Each fragment finishes + before the next resolution can evict its weights. Batches above the configured + limit retain TurboMind. The opt-in also affects prefill tails of 21–40 rows + and therefore requires numerical and model-quality evidence. +* `qsa_host_kv_state_compact` defaults to false. The bounded host-KV allocator + gives each recurrent group independent device banks and distinct scheduler ID + ranges. The block tables and zeroer translate those IDs into local bank IDs; + null ID zero remains zero. History IDs, prefix hashes and reference lifetimes + remain global. PP or CP above one is explicitly rejected. This avoids charging + every layer for every recurrent group's IDs. For C8, the tested layout saves + 463.46 MiB per rank, rather than the initially estimated 1.6 GiB. +* `sm70_qsa_shared_key_batch` defaults to false. It reuses the existing + single-request native scorer for each request in a small batch, then keeps + shared batched top-k/expansion. Admission requires FP16 H4/D128 on SM70, + two to eight queries per request, valid CPU query starts and at most 64 rows. + All requests are admitted before any scorer executes; other layouts retain + the paged scorer. No new native provider or kernel is introduced. +* `sm70_qsa_indexer_single_rank` is an experimental opt-in, default false. + It binds only for SM70 FP16 H4/KV1/D128, host compressed-indexer storage and + TP4/PP1/DCP1/PCP1. With one to 64 rows, rank zero selects indices and the + existing TP communicator broadcasts them. Every rank still writes its own + raw/compressed keys. Larger batches retain independent selection. Device + indexer history, including the qualified C1 layout, remains on its existing + path. The 32K task gate and fresh ordinary-worker paired task/performance gates pass (see the updated evidence below). + +These are explicit opt-ins until the paired end-to-end gates pass. No kernel +sidecars, copied extensions or preloads are required. + +## Validation and negative results + +The five QSA GPU cases cover M32, M33, M40, M64 and M65, including changed KV +pages, selection indices and lengths under graph replay. They pass against +the protected reference reader. The compact pool tests cover C1/C2/C4/C6/C8, +the adjacent insufficient capacity, unique per-group quotas, null IDs and +prefix-cache lifetime. CPU cache/dispatch tests pass (46 cases), and typed +policy/CPU block-table tests pass (12 cases). + +Four-card GPU tests pass (23 cases): compact-state zeroing with guarded foreign +IDs and changed graph inputs; block-table request permutation and padding; +tiled GGUF experts against independently dequantized Q8 oracles at M20/M21/M30/M40; +and compiled HCX payload ownership at M5/M8/M9/M10/M20/M30/M40/M41, including +changed inputs and final mixing fallback. GGUF error bounds are rtol/atol 0.003; +HCX retains the existing 0.002 relative-norm allowance. These operator tests +do not establish full-model quality or performance. + +Grouped scoring passes 14 GPU cases with device and mapped-host histories, +C2/C4/C6/C8, request reorder, missing pages, zero lengths and padded rows. +Changed-input graph replay uses rtol 0.003/atol 0.0002 for scorer reductions; +expanded selection matches the independently invoked native request scorer +exactly in the tested cases. The single-rank prototype passes a four-rank +graph test at M1/M5/M10/M20/M30/M40/M64/M65: only rank zero executes selection +through M64, while M65 falls back on all ranks. Changed requests/pages/lengths +match independent selection exactly. Typed policy/hash tests pass too. + +An aligned main-history cache microbenchmark rejected increasing hot capacity +from 32768 to 49152: both are about 0.219 ms for the representative expanded +indices. An earlier unaligned synthetic selection produced about 10 ms, but +does not represent the indexer's four-token expansion and is not model evidence. + +Combined synchronized worker events and Nsight node tracing failed twice +(one CUDA launch/Xid43, one CUDA-event recording segmentation fault during +mixed prefill). All five event-only measurement batches completed before the +second trace failure. Subsequent event and node passes are separate; partial +traces are used only for attribution. Those failures are not discarded or +reported as successful full-width trace coverage. + +The first compact C8 start captured M5/M10/M20/M30/M40 but its hand-built +startup warmup allocated IDs outside the new state-group ranges. Warmup now +uses the existing BlockPool and respects each group's capacity; 17 warmup +tests pass. The next start exposed an active-M33/padded-M40 output gate shape +error in the existing host QSA wrapper. Active gates are now sliced with the +other active inputs; six CPU wrapper cases cover empty input, M3/M5, M7/M10 +and M33/M40 padding. These failed starts remain separate from performance data. + +Rejected unmodified C8 starts include 1024-token and 512-token chunks that ran +out of device memory. A 64-token chunk with M30 graphs instead entered +TurboMind prefill and attempted an event dependency on an uncaptured stream. +This is distinct from CUDA IPC handle failures; #1208 is not a prerequisite +for that reproducer. A later M5-only C8 start also ran out of memory. Raw failed +logs and contracts are preserved; none is a speed baseline. + +## Reproduction and retained artifacts + +Run the service with the qualified sampling/MTP/KV settings and an owned GPU +group. Use `VLLM_SM70_MTP_PROFILE=1` and +`VLLM_SM70_MTP_PROFILE_INTERVAL=1` only for the diagnostic phase pass. +The benchmark requires a JSON list of prompt cases with `name` and `prompt`. + +```bash +python benchmarks/flashnext_prefill/measure_mtp_concurrency.py \ + --base-url http://127.0.0.1:18041 --cases /path/to/prompt-cases.json \ + --concurrency 1 2 4 6 8 --service-log /path/to/worker.log \ + --out /path/to/requests.json +python benchmarks/flashnext_prefill/analyze_mtp_verifier.py \ + /path/to/worker.log /path/to/requests.json /path/to/verifier.json +``` + +For node attribution, launch the same service under Nsight with +`--cuda-graph-trace=node:host-only --capture-range=cudaProfilerApi` and +`--profiler-config '{"profiler":"cuda"}'`. The optional diagnostic worker +`mtp_trace_worker.TraceWorker` annotates target and draft graph submissions +after graph capture. Capture only short warmed generation intervals. Report +GPU service sums, same-rank verifier wall and TP4 envelope separately; node +tracing is composition evidence, not an accepted absolute latency. + +Local raw artifacts are retained at +`/home/ymzx/桌面/1cat-vllm/reports/flashnext-mtp-batch-20261014/`. +The local public API is stopped. The 9683 API is outside this experiment and +remains running. No shared environment packages or model files are modified. + +## Single-rank host-indexer candidate + +Only the single-rank selection option changes relative to the full-width +grouped-scoring event pass. All five measured request streams complete; +the collector performs no Nsight capture. Values remain synchronized worker +diagnostics, not accepted uninstrumented throughput. + +| Concurrency | Complete rounds | Before ms | Candidate mean ms | p50 ms | p90 ms | p99 ms | Forward ms | Sampling ms | State ms | Draft ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 27.422 | 20.979 | 20.823 | 21.678 | 23.439 | 20.557 | 0.404 | 0.017 | 4.124 | +| 2 | 197 | 52.964 | 38.099 | 37.782 | 38.939 | 43.372 | 37.648 | 0.432 | 0.019 | 5.330 | +| 4 | 169 | 93.346 | 58.831 | 58.283 | 61.249 | 65.323 | 58.283 | 0.525 | 0.023 | 7.673 | +| 6 | 168 | 138.721 | 84.812 | 84.353 | 86.517 | 91.666 | 84.198 | 0.588 | 0.026 | 10.236 | +| 8 | 178 | 185.889 | 111.006 | 110.241 | 113.423 | 123.378 | 110.115 | 0.861 | 0.030 | 12.568 | + +C2 improves a further 28.1%, C4 37.0%, C6 38.9% and C8 40.3%. +C8 is 6.94 times the approximately 16 ms device-indexer C1 reference; +its 13.876 ms per-request verifier share is not a 13.876 ms batch wall. +The candidate host-indexer C1 is 20.979 ms, so it does not meet the qualified +device-indexer C1 latency. Draft costs are separate. Eight submitted 32K retrieval/arithmetic requests, executable order-preserving +Python deduplication, automatic weather-tool calls and tool replies pass with +thinking, production sampling and natural EOS. The tool test initially required +`finish_reason=stop` for a successful tool call; that harness mistake was fixed +to admit the expected `tool_calls` completion, and tool/reply checks were rerun. +Fresh ordinary-worker paired task and emitted-token measurements completed; see the updated evidence below. Further batch latency optimization remains necessary. +The affected final CPU suite passes 73 cases. A stale warmup fixture originally +patched the removed runner platform/environment hooks; it now exercises the +existing bound warmup policy. No runtime compatibility workaround was added. + +## Updated acceptance target and completed clean control + +The requested costs refer to the **whole MTP4 verifier batch**, not its share +per request: C1 approximately 16 ms, C2 17–18 ms, C4 approximately 22 ms and +C8 approximately 28 ms. The previous 111 ms C8 result does not satisfy that +objective even when divided by eight. All further results retain five target +rows and four draft rows per request, and separate draft time. + +Fresh ordinary workers with profiling and NVTX disabled completed paired +root-selection controls, using the same capacity-eight host-indexer setup. +Common decode intervals exclude prefill and generation boundaries: + +| Concurrency | Retained independent selection tok/s | Root selection tok/s | +| ---: | ---: | ---: | +| 1 | 76.46 | 107.91 | +| 2 | 99.35 | 116.72 | +| 4 | 94.74 | 157.36 | +| 6 | 95.92 | 144.65 | +| 8 | 98.71 | 150.58 | + +Both runs pass the same task gates: eight seeded 32K exact-field retrieval +answers with natural EOS, executable deduplication code, automatic weather +tool invocation and its response. These are emitted-token speed evidence, +not replacement measurements for the verifier cost. The local GPU-indexer +C1 control is 14.655 ms and C2 25.472 ms at capacity two. A capacity-four +GPU-indexer experiment failed near full VRAM and is not accepted speed evidence. + +## Bounded shared-page indexer staging + +`sm70_qsa_indexer_stage_keys` defaults to false. A qualified host FP16 +H4/KV1/D128 indexer borrows one engine-owned device bank, bounded to 32768 +compressed keys rounded to physical pages (about 8 MiB at page size 204). +Every selection resets its map, deduplicates live physical pages across up to +eight requests, and refreshes their keys. The host history stays authoritative; +missing pages and overflow retain the existing scorer's paged-reader semantics. +The native `run_staged` entry reuses its original arithmetic; `run` keeps its ABI. +PP/DCP/PCP above one, other query layouts and absent native entries retain their +qualified paths. The bank never exposes borrowed output views and does not +cross engine configuration transfers. + +The short scorer test at 32K improves C8 shared-prefix scoring from 3.74 to +0.36 ms; independent histories improve approximately 1.9x. This is not a claim +that arbitrary independent contexts obtain the shared-prefix gain. Twenty-three +GPU tests pass, including changed key contents, page reorder, padding, missing +pages, graph replay, exact selection and a deliberately overflowing one-page bank. + +The full model diagnostic uses an 8 MiB QSA score tile and the same sampler, +256K capacity, prefix cache, 80% GPU / 20% host experts and MTP4. A previous +32 MiB score-tile run exhausted VRAM on a 64 MiB prefill score allocation; +that failure is retained. The bounded setting completes all 32K shared-prefix +streams with no incomplete TP4 verifier rounds: + +| Concurrency | Rounds | Whole verifier mean ms | p50 ms | p90 ms | p99 ms | Draft mean ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 17.887 | 17.733 | 18.616 | 19.950 | 3.881 | +| 2 | 193 | 29.413 | 29.124 | 30.535 | 32.590 | 4.610 | +| 4 | 191 | 36.992 | 36.497 | 38.333 | 41.594 | 5.935 | +| 6 | 181 | 52.411 | 51.824 | 54.096 | 60.300 | 7.513 | +| 8 | 178 | 66.309 | 65.201 | 69.296 | 80.457 | 8.807 | + +These are diagnostic GPU events with critical-rank alignment, not clean +production throughput. They remain above the requested whole-batch costs. +The staging candidate's fresh clean paired throughput, long-output drift and +256K boundary gates remain pending; no default is promoted. + +## Follow-up candidates and rejected directions + +Large dense projections now have an opt-in `sm70_gguf.small_m_hmma_batch`. +It reuses the existing segment format and native decoder across token tiles; +qualified DMV Q4/IQ4 banks provide their equivalent full scale representation. +The six-array prepared-operand ABI is retained. The accelerated initialization +subset is K2560/N>=1536 or K1536/N2560 in the existing qualified formats and +layouts. Runtime acceleration covers M9..16 and M25..40; M1..8 and +M17..24 retain the previous projection routes. Compatible small segment +peers share the same batch launch. Larger batches retain canonical restoration. The native batch operator has separately tested M1..40 coverage. The old native entry retains +its M1..32 contract; the batch entry is explicit. Eighty-six dense native GPU +cases and 111 affected storage/runtime cases pass, including capture/replay +and the M40/M41 fallback boundary. The initial wider M9..40 model experiment +produced C1/C2/C4/C6/C8 costs 17.931/26.655/38.265/51.922/63.368 ms: +C2 improves, but C4 regresses. That wider runtime admission is rejected; +M9..16 retains its C2 opportunity and preserves the prior M20 route. The +narrowed policy still needs matched full-model evidence. + +A single-grid HC experiment passed its small numerical check but did not +improve C8 (approximately 0.192 ms for both retained and candidate calls). +Its implementation is removed from the runtime and retained as a research +patch in the private artifacts. Existing eight-row HC tiles remain qualified. + +A separate default-off `sm70_gguf.grouped_gate_up` candidate reuses the +original lattice weight decoder for up to four routes to one expert. The +accelerated shape is M9..20, source types18/21/22, N divisible by32 and K by256, +with the existing Q8_1 intermediate. Routing is sorted and admitted on the GPU +each replay; outputs return directly to their original route order. Sparse +routes retain the original single-route algorithm. Changed-input graph tests +exercise both live admission outcomes and exact Q8_1 bytes; full-model benefit +remains unqualified. The first merged grouped/sparse kernel regressed sparse +routes because they inherited its register and shared-memory footprint; it is +rejected in favor of a separately guarded retained kernel. + +Rejected experiments are preserved in the task artifacts: + +* Parallel mapped-host indexer scoring in request-major and interleaved grids + slowed independent-history scoring approximately threefold. They are removed. +* The retained full-batch HC matrix path is slower at C1/C2 and improves C8 + only slightly; it is not promoted as a scaling fix. +* Expert lane-width and bank-layout changes improve the leading integer kernels + by only about 10%; they do not establish the requested scaling. +* Upstream MoE MMQ at M40 takes 1.6–1.9 ms for gate/up in representative real + banks, versus about 0.22 ms for the retained integer route; it is rejected. +* Sorting routed experts alone passes exact Q8 output checks but slows the + complete scorer by approximately 2–4%; its private native entry is removed. +* The expert-cache entry already holds 248 route IDs. At top-k ten the C8 + native path uses two twenty-row fragments. The 136 count in startup logs is + host weight slots, not request-entry capacity; it does not justify a ring fix. + +Hardware counter attribution is not available to the local unprivileged user +(`RmProfilingAdminOnly=1`). Cache/decode hypotheses are therefore based on +paired exact-shape experiments, not claimed Nsight Compute counters. + +## Combined narrowed dense and grouped-expert diagnostic + +The narrowed dense policy (M9..16) plus grouped experts completes all five +32K shared-prefix cases under the same capacity-eight, score-tile 8 MiB, +chunk 512, host-history contract. Whole verifier includes target forward, +sampling and state commit; draft remains separate. + +| Concurrency | Complete rounds | Whole verifier mean ms | p99 ms | Draft mean ms | +| ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 17.768 | 20.321 | 3.871 | +| 2 | 202 | 26.779 | 31.678 | 4.601 | +| 4 | 203 | 36.705 | 41.995 | 5.915 | +| 6 | 203 | 51.769 | 59.463 | 7.472 | +| 8 | 190 | 64.895 | 80.106 | 8.808 | + +There are no incomplete TP4 rounds. This combined experiment does not isolate +the contribution of grouping: the preceding wider dense policy also differs. +Relative to page staging alone, its C8 improvement is about 2.1%; this is not +the requested 28 ms C8 target. At C8 forward is 64.005 ms; sampling 0.860 ms and +state commit 0.030 ms cannot explain the inflation. A new matched node trace +will decide the next bottleneck rather than attributing it from micro gains. + +The eight submitted 32K retrieval/arithmetic tasks, executable Python function, +automatic weather-tool call and its reply all pass with natural EOS. The +quality run still has worker diagnostic fences, so clean emitted-token +qualification and the 256K boundary remain pending. The first model run uses +the owned native build immediately before a symbol rename/formatting pass; +the final formatted source artifacts are subsequently rebuilt and installed +atomically. The final artifacts pass 54 affected GPU cases, including changed-route +graph replay, independent expert output oracles and four-rank indexer state +checks. Clean model speed and longer-context gates remain separate before +any promotion. + +## Matched node attribution and mixed-segment fusion + +Five decode-only node captures use the final rebuilt source and identical +C1/C2/C4/C6/C8 contracts. Their interior counts are150/113/87/60/49. +Node tracing changes absolute latency: traced whole verifier means are +22.105/29.934/39.397/57.967/71.533ms, separately from event measurements. + +| TP0 GPU service ms | C1 | C2 | C4 | C6 | C8 | +| --- | ---: | ---: | ---: | ---: | ---: | +| Dense/shared projections | 3.489 | 6.879 | 6.224 | 9.291 | 14.847 | +| Routed expert projections | 2.880 | 4.639 | 7.573 | 12.015 | 15.439 | +| HC and its fused exchange | 2.504 | 4.648 | 7.564 | 10.575 | 13.877 | +| Indexer and top-k | 3.556 | 3.876 | 4.434 | 5.109 | 5.955 | +| Expert-cache resolution/wait | 1.659 | 1.398 | 1.498 | 3.285 | 3.761 | + +These are service sums, with overlapping streams; they are not an additive +wall decomposition. C8 peer broadcast includes waiting for TP0 indexer work. +The actual local GPU4–7 mesh has mixed NV1/NV2 links; the remote9683 NV2 mesh +is a separate hardware contract. CPU affinity matches NUMA node1. + +The captured operand descriptors expose a dispatch problem: QKV banks mix a +large batch-enabled segment (family5) with existing small K/V segments +(family3). Requiring every segment to have a batch tag splits compatible +projections and restores the small peers before separate calls. The candidate +admits one native batch launch when all segments have the compatible layout +and at least one calibrated large segment enables batching. The old M20 +fusion remains unchanged. M9..16 and M25..40 are qualified separately from +larger prefill fallback. + +Real-weight QKV screens across four layers improve M10 by2.07–3.63x, +M30 by1.48–2.04x and M40 by1.98–2.81x. M20 arithmetic is unchanged. +Relative output errors remain below0.00077 versus the FP32 dequantized-weight +reference, under the predeclared0.003 norm bound. The affected111 GPU runtime +cases pass; three additional mixed-format cases exercise11 boundary/interior +M values, changed-input replay and exactly one native batch dispatch at the +accelerated sizes. The matching complete model curve is being collected; +these micro gains are not reported as whole-verifier gains. + +Two original-down scheduling screens are rejected: cooperative two/four-lane +source reads usually regress despite exact outputs, and one-warp-per-route +regresses IQ4_NL while offering no representative M20 gain for Q2_0. Both +research patches and all15 real-weight cases are retained outside the runtime. + +The mixed-segment model experiment completes all five 32K shared-prefix +cases with181/214/182/178/165 complete TP4 verifier rounds and no incomplete +rounds. Whole means are17.513/26.114/36.692/50.150/61.131ms; C8 p99 is68.697ms +and draft mean8.768ms. C8 improves5.8% relative to the preceding combined +64.895ms diagnostic. Retrieval, executable code and automatic tool call/reply +all pass. This remains above the requested18/22/28ms targets and uses worker +fences; ordinary-worker qualification remains separate. Source/native sampler +acceptance varies between runs, so HTTP gains alone are not attributed to the +projection change. + +An isolated HC grid-barrier screen compares the existing cumulative arrival +counter with distributed arrival flags. Both take approximately1.62us per +barrier and pass ordering/wrap checks in that screen. No native protocol change +is justified; the private prototype is retained as a negative result. + +## Closed verifier interval audit + +The five captures use the pre-mixed-fusion grouped-twenty-row candidate +(the untraced C8 event mean was 64.895 ms). They were swept by GPU timestamp on each +selected critical rank. Each interval belongs to exactly one category, a +separate overlap bucket, or the idle residual. Categories below sum to verifier +wall; they are not summed kernel service. Node tracing perturbs absolute latency +and these numbers must not replace the untraced event measurements above. + +| Concurrency | Complete interior rounds | Verifier wall ms | HC ms | Experts ms | Dense/shared ms | TP collectives ms | Overlap ms | Idle ms | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| 1 | 150 | 22.105 | 2.390 | 2.879 | 2.489 | 3.661 | 0.999 | 2.150 | +| 2 | 113 | 29.934 | 4.592 | 2.411 | 3.577 | 3.988 | 5.689 | 2.590 | +| 4 | 87 | 39.397 | 7.399 | 7.560 | 5.488 | 4.823 | 0.167 | 3.035 | +| 6 | 60 | 57.967 | 10.565 | 11.416 | 5.086 | 5.821 | 4.334 | 4.364 | +| 8 | 49 | 71.533 | 13.702 | 12.595 | 7.753 | 6.724 | 9.138 | 4.338 | + +The remaining wall belongs to selected-history attention, GDN, routing, cache +resolution, normalization, sampling/state commit and uncategorized kernels. +Raw `scaling-nodes-{1..5}-closed.json` files retain those categories and each +round; `scaling-nodes-closed-summary.json` retains their means. The closure +assertion tolerates at most 1e-7 ms. This audit includes target forward, +sampling and state commit, and excludes MTP draft calls. Expert and HC scaling +remain primary optimization targets. A collective can include a wait on the +indexer root, so its duration does not establish intrinsic link cost. + +## Forty-row expert/cache experiment + +The default fragment and mailbox capacity remain twenty rows. An explicit +`sm70_gguf.grouped_gate_up_max_m=40` permits one 400-route lease with a +512-integer mailbox. Reply capacity accounts for all request entries plus +spare restoration; it grows to the next power of two when required. Retained +quantization and ordinary gate/up/down kernels still run in twenty-row tiles. +Grouped original gate/up has a separate native forty-row entry point, keeping +the legacy operator and opaque schemas unchanged. + +The matched diagnostic event curve with forty-row gate/up, without grouped +down, is 17.752 / 26.270 / 37.187 / 47.641 / 61.115 ms for C1/2/4/6/8, +with 202 / 200 / 190 / 196 / 174 complete rounds and no incomplete rounds. +C6 improves versus the mixed-fusion twenty-row control; C8 does not. This is +not evidence that a faster isolated gate/up microbenchmark speeds the service. +Eight concurrent 32K exact retrieval/arithmetic responses, executable Python, +automatic weather tool selection/reply and natural EOS checks all pass. + +The next default-off `sm70_gguf.grouped_down` candidate groups original +IQ4_NL/Q2_0 down weights across at most four equal-expert routes. It keeps +FP16 projection boundaries, the retained four warp-owned route subtotals and +their addition order. Only M25..40 is accelerated, with the forty-row lease +policy and admitted Q8 intermediate/native expert shape. Smaller fragments +keep the retained down path. Unsupported shapes use that path or its existing +TurboMind continuation; a requested qualified route without its native build +rejects initialization. + +Twenty-seven real-weight down screens and twenty-one threshold/interior +screens preserve the prior bytes, including the saved prior-binary references. +At M40, sort-inclusive speedups are 1.12..1.92x for IQ4_NL and 1.17..1.57x +for Q2_0. Sparse M20 IQ4_NL regresses and remains outside acceleration. Sparse +M25 IQ4_NL only improves 1.013x; this boundary result is not a broad speed claim. +The native, cache, state/graph and complete expert-chain checks pass 141 GPU +cases after correcting a test parameter that shadowed its weight tensor. +The combined service screen below establishes measured 32K performance and +task evidence; 256K quality qualification remains outstanding. + +Additional rejected screens are retained in the report directory. A bounded +mailbox-plan scan changed isolated cache time by less than 1% and was removed. +A four-float HC partial-workspace pad appeared faster in a one-call graph, +but the 64-call graph screen removed launch skew and showed no C8 gain +(0.13680 ms retained versus 0.13803 ms padded). Its source was removed. + +## Rechecked whole-batch HC path + +The earlier whole-batch HC rejection did not establish a reliable native-path +comparison: its launcher did not explicitly enable the fusion policy, and a +single-call graph mixed GPU work with peer launch skew. Rechecking the native +operator directly, with observed CUDA dispatch and sixty-four calls per graph, +shows bit-identical outputs and critical-rank per-call times: + +| Concurrency | Retained tiled HC ms | Whole-batch HC ms | Speedup | +| --- | --- | --- | --- | +| 1 | 0.021842 | 0.021777 | 1.003x (retained fallback) | +| 2 | 0.045594 | 0.043849 | 1.040x | +| 4 | 0.073743 | 0.065897 | 1.119x | +| 6 | 0.103555 | 0.092633 | 1.118x | +| 8 | 0.136066 | 0.121190 | 1.123x | + +This supersedes the earlier HC rejection for the measured full-mesh TP4, +separate-projection, local-schedule path. The default-off typed policy +`sm70_hcx_batch_fusion` owns its larger peer buffers and local workspaces. +Producer-fused output projections, non-full-mesh execution, unsupported local +modes and M1..16 retain the original eight-row path. M9..16 was excluded after +the ordinary-worker acceptance regression described below. M>40 retains the existing +projection/collective continuation. Requested qualified fusion rejects a +missing native build. Original operator signatures remain unchanged. + +The four-rank Inductor/graph payload test passes all M5/8/9/10/17/20/24/25/30/ +32/33/39/40/41 cases, including input mutation, retained/fast alternation, +materialized/final HC consumers and output ownership checks. These results are +kernel/state evidence. The combined service performance and task gates are +being measured separately; 256K boundary qualification remains outstanding. + +## Combined expert-down and HC service screen + +With mixed dense fusion, forty-row expert leases, grouped down and whole-batch +HC enabled, the matched MTP4 diagnostic GPU-event verifier costs are: + +| Concurrency | Complete rounds | Incomplete | Mean ms | P99 ms | Draft mean ms | +| --- | --- | --- | --- | --- | --- | +| 1 | 202 | 0 | 17.650 | 20.388 | 3.875 | +| 2 | 220 | 0 | 25.681 | 28.084 | 4.580 | +| 4 | 214 | 0 | 35.286 | 42.640 | 5.936 | +| 6 | 186 | 0 | 45.974 | 58.584 | 7.455 | +| 8 | 193 | 0 | 54.660 | 64.421 | 8.749 | + +C8 improves from 61.131 to 54.660 ms (10.6%) versus the prior mixed-fusion +screen. C4 improves from 36.692 to 35.286 ms and C6 from 50.150 to 45.974 ms. +These are whole-batch costs, never costs divided by concurrency. The workload +keeps 256K capacity, 32K shared input, MTP4, C8 capacity, the same host QSA and +80% GPU expert storage, 512-token prefill chunks and natural-EOS sampling. +Profiling events/fences are enabled. The 16/18/22/28 ms C1/2/4/8 targets remain +unmet. Ordinary-worker throughput/acceptance and paired task checks are a +separate gate; no default promotion follows from this screen. + +## Ordinary-worker acceptance audit + +Fresh ordinary workers, without event fences, NVTX tracing or worker wrappers, +completed matched mixed-fusion control and combined expert-down/HC passes. +Both use the same source/native artifact, 384 output tokens, 32774 input +tokens, capacity eight, sampling and per-request seeds. The control retains +twenty-row gate/up and tiled HC; the combined pass enables forty-row gate/up, +grouped down and whole-batch HC. Each row is the measured common interval +with all submitted requests still decoding; rates are aggregate emitted tokens. + +| Concurrency | Control tok/s | Combined tok/s | Control acceptance length | Combined acceptance length | +| --- | ---: | ---: | ---: | ---: | +| 1 | 110.15 | 125.34 | 2.477 | 2.824 | +| 2 | 169.49 | 133.84 | 2.630 | 2.121 | +| 4 | 213.31 | 235.15 | 2.339 | 2.473 | +| 6 | 251.40 | 277.44 | 2.458 | 2.481 | +| 8 | 259.09 | 305.02 | 2.259 | 2.451 | + +The C2 emitted-token regression is approximately 21%, despite a slightly +faster diagnostic verifier. This candidate is not accepted as a general MTP +speed improvement. C1 acceptance also changes, so its rate improvement cannot +be attributed to verifier compute alone. Both passes complete the eight-way +32K retrieval/arithmetic, executable Python, automatic tool-call/reply and +natural-EOS gates. These bounded task checks do not establish unchanged MTP +acceptance or broad quality. The runtime HC acceleration range is consequently +narrowed to M17..40 while C2 acceptance is investigated. The native operator's +correctness range remains M9..40; admission and measured benefit are distinct. + +The subsequent shared-expert experiment reuses the existing fused gate/up and +down operations in bounded twenty-row fragments. It is explicitly enabled by +`sm70_gguf.shared_expert_batch` during the experiment, admitting Q8_0 down projections at +M9..40, and preserves the canonical continuation for other projections and +larger batches. The existing two-integer internal down descriptor remains +valid; the explicit opt-in adds a third policy integer without changing the +registered operator schema. Real-weight screens across six layers show +1.34..2.18x isolated M40 speedups and paired relative-norm error at most +0.000622 against the predeclared 0.003 bound. M5/M8/M20 retain their prior +bytes. GPU graph tests cover M5/8/9/10/16/17/20/21/24/25/30/32/33/39/40/41, +mixed source formats, the independent FP32 oracle, changed inputs and native +execution/fallback. Eleven format combinations pass. + +The first ordinary-worker shared-expert pass emits 110.24 / 151.25 / 244.01 / +253.33 / 292.45 tok/s at C1/2/4/6/8. C2's additional same-seed repetition emits +159.54 tok/s, with acceptance length changing from 2.444 to 2.579. Eight-way +retrieval/arithmetic, coding and tool/reply gates pass. C6/C8 do not exceed the +previous combined pass, and C2 remains below control. This is an unpromoted +experiment; an isolated shared-expert speedup is insufficient to accept it. +Its separate synchronized-verifier curve is 17.701 / 26.318 / 36.261 / 47.069 / +61.686 ms, with 181 / 192 / 181 / 181 / 180 complete TP4 rounds and no +incomplete rounds. The experiment does not improve whole-verifier cost and +was removed, including its policy and runtime code. Its patch and tests remain +in the raw artifact directory; operator micro speed does not justify retaining +another runtime option. + +A subsequent ordinary run keeps shared-expert acceleration disabled and HC +acceleration limited to M17..40. C2 repetitions emit 142.16 / 150.06 / 155.10 +tok/s with acceptance length 2.326 / 2.326 / 2.568. Its C1 cold prime and warm +repeat have acceptance length 3.031; the eight short math primes used by the +earlier paired run were omitted. Thus this isolated run does not establish +recovery or isolate a numerical cause. Per-request seed alone is insufficient +to treat these different runs as identical generation histories. The HC +admission restriction remains while the acceptance variance is unresolved. + +The latest combined C8 node capture closes a 60.184 ms traced verifier wall +with 12.465 ms HC/fused communication, 9.069 ms expert projections, 6.599 ms +dense/shared projections, 6.453 ms TP collectives, 3.874 ms selected-history +attention, 2.423 ms GDN, 9.191 ms overlap, 3.399 ms idle and 6.711 ms in the +remaining categories. Categories are exclusive wall intervals; traced absolute +time is not the 54.660 ms diagnostic result. The observed native expert route +groups average approximately four rows per expert. Neither these counts nor +collective durations establish memory bandwidth or intrinsic NVLink cost; +Nsight Compute profiling counters remain unavailable to this user. + +## HC scheduling and partial-write follow-up + +The instrumented HC screen partitions each CTA's control-thread interval, +separately from the ordinary-worker verifier measurement. At C8 its mean +117.942 us consists of 24.627 us initial peer writes, 4.472 us weight prefetch, +26.112 us peer reads/combine, 14.968 us down partials, 5.750 us first grid +rendezvous, 18.591 us RMS/down reduction, 7.816 us LoRA peer writes, +0.466 us LoRA completion, 3.310 us second grid rendezvous, 5.949 us LoRA +shared-memory load, 5.164 us HC up/output writes and 0.717 us output completion. +These are instrumented per-CTA phase intervals, not hardware counters or a +sum of per-stage rank maxima. + +The retained small-HC local schedule issues three independent peer loads +before inspecting their epochs. The first whole-batch kernel instead waited +for each peer in sequence. The candidate reuses the retained independent-load +schedule with the same tags, timeout and canonical addition order. Its down +projection transposes partials in shared memory before coalesced global writes. +The LoRA input tile is unused during this phase and owns that temporary view; +no new shared-memory allocation, host pool or function attribute is required. + +The private research screen, using 64 calls per graph and twenty timed replays +on all four ranks, improves C8 HC from 0.121 to 0.104 ms. Outputs match both the +packaged native baseline and private retained kernel byte-for-byte; changed +graph inputs pass too. This screen only admits the idea to the ordinary native +build. Fresh normal-artifact graph/state and model measurements remain the +qualification gate; its private research library is not a service dependency +or a production speed result. + +The normal CMake artifact subsequently passes 44 CPU policy/dispatch cases, +the four-rank compiled payload/state suite and the direct native changed-input +graph screen. C1/2/4/6/8 direct native costs are 0.022427 / 0.042894 / 0.059882 / +0.083102 / 0.107008 ms; C1 is a retained fallback. Corresponding retained +eight-row tiles are 0.022294 / 0.045446 / 0.075752 / 0.104005 / 0.136468 ms. +The native correctness range includes M10; runtime acceleration still begins +at M17. The new ordinary `_C` SHA256 is +`283b58ff6a99f358b74a4c2d579f6757b996cbe4ca79f908c721cbf5e03cda77`. +Whole-model qualification remains separate. + +A subsequent two-shape expert-down screen compacts four-route group leaders +and schedules persistent work blocks, retaining the original integer dot and +unroute order. At M40, 128 expert IDs and 149 active groups, the sort/plan-inclusive +cost changes 0.102674 to 0.100932 ms for IQ4_NL and 0.077722 to 0.076186 ms for +Q2_0; outputs remain byte-exact. These 1.7%/2.0% gains fail the predeclared +8% screen threshold. The research source is archived and not integrated. + +## Latest whole-model HC result and next indexer candidate + +The normal coalesced HC artifact completes ordinary task checks and the +separate synchronized-verifier run. Whole-batch verifier means at C1/2/4/6/8 +are 17.652 / 26.573 / 36.066 / 43.682 / 53.886 ms, with 187 / 202 / 222 / +189 / 199 complete TP4 rounds and no incomplete rounds. Draft means are +3.857 / 4.604 / 5.949 / 7.499 / 8.754 ms. These are whole batches, without +dividing by request count. C2/C4/C8 targets of 18/22/28 ms remain unmet. +Ordinary common-interval aggregate rates are 125.41 / 138.15 / 239.56 / +261.18 / 298.48 tok/s; additional C2/C8 repetitions give 141.43/327.25. +Retrieval/arithmetic, executable code, automatic tools and natural EOS pass. +The C2 regression remains unresolved, and this is not a general MTP speed win. + +A two-hop HC peer reduction retains bytes but regresses the direct graph +screen: C4 0.061217 to 0.071124 ms and C8 0.105724 to 0.126098 ms. It remains +research-only and is rejected. Its extra rendezvous does not amortize the +smaller peer fanout. + +The next default-off candidate is `sm70_qsa_indexer_persistent_keys`. It +reuses the existing FP16 shared-key scoring family, with one bounded per-layer +bank plus the existing shared refresh bank. Physical-page tags and int64 +write versions identify reusable pages. Both fused and reference compressed +writers advance versions before selection or its MTP skip path. Versions +never rewind with speculative visibility. Every cache bind clears tags, +including same-tensor rebinding before recapture. Hash collisions use freshly +copied spill pages, then the existing authoritative paged reader for overflow. +The default staging path still refreshes every live page on every selection. + +The accelerated subset is the existing SM70 FP16 H4/KV1/D128 host-indexer +path with TP4 root selection, CP1 and PP1. External KV transfer, KV offload +and speculation methods other than ordinary MTP have untracked writers and +are explicitly rejected when this candidate would be selected. Other layouts +retain the existing path. A mismatched cache binding retains protected +refresh; closed owners reject reuse. Normal native loading requires the +matching `run_staged_banked` operation; prior native schemas remain unchanged. +This changes neither attention selection arithmetic nor KV storage format. + +At compressed page size204, 16384-key per-layer banks add approximately52.4MiB +on the selection root across twelve target layers and one MTP layer; actual +worker snapshots confirm thirteen persistent banks. The shared refresh bank +remains32768 keys. Startup, cold prefill and captured C8 generation must +fit the unchanged expert placement and memory contract. A private NUMA1 +operator screen improves shared-prefix C8 from0.367 to0.197ms and independent +C8 from2.976 to2.309ms. These local figures do not establish whole-verifier +speed, actual256K performance, broad numerical coverage or production quality. +The one-bank variant regressed independent C8 by7.6% and was rejected before +runtime integration. Normal-artifact state and model gates remain separate. + +The normal banked-indexer artifact subsequently completes the model gates. +Its whole-batch C1/2/4/6/8 verifier means are15.925 /24.678 /33.724 /41.812 / +53.130ms, with187 /181 /201 /213 /193 complete TP4 rounds and zero incomplete +rounds. Corresponding p99 values are18.019 /28.322 /46.876 /59.045 /66.672ms. +The single-stream mean returns to approximately16ms under this capacity-eight +contract, while the desired C2/C4/C8 means remain unmet. C8 improves only0.756ms +in this pass; its isolated indexer improvement is not a whole-batch speed claim. + +Fresh ordinary workers emit122.06 /182.21 /251.03 /288.93 /322.03 aggregate +tok/s; extra C2/C8 repetitions give169.11 /308.37. Acceptance lengths change +and C8 overlaps the prior repetition range. All five retrieval/arithmetic, +executable-code and tool/reply gates pass, as do eight short natural-EOS checks. +Observed peak root GPU allocation is16078MiB, leaving67MiB available on the +16384MiB card after driver reservations; other ranks use16050..16052MiB. +This fits the measured32K workload with unchanged80% GPU expert placement, +but does not establish256K boundary memory or quality. No default is promoted. + +The normal indexer SHA256 is +`ca41596598165c7ecb39c127f4b6c82ef5458dc5bdad16cf4abc4ca8e1237842`. +Fresh-process dependency inspection and all four actual worker maps confirm +ordinary source-owned extensions without a private research DSO or preload. +Actual worker JIT-monitor entries confirm version tracking and bank planning. +The first48 GPU and29 policy/cache tests, six writer-ordering cases and full +changed-source hooks pass. A further18 GPU cases include the actual204-token +compressed page size, stale-binding refresh fallback, same-address generation +reset and changed-input replay after cache rebinding and recapture. Hash +collision winners may change between launches; zero-copy reuse is required +for unchanged pages with stable distinct leases, not for conflicting hashes. + +A follow-up dense screen tests smaller row chunks in the existing weight-major +family. The trace shows244 registers/thread,112..128 CTAs of128 threads for +large input projections, and80 CTAs of256 threads for the1536-to2560 output +projection. These are launch resources, not measured occupancy or stall +counters. Smaller chunks reduce registers to approximately168 but reload and +decode weights. Across six M20/M40 shape screens, the dominant Q6_K input +projection does not improve. One mixed IQ4_XS/Q8 case improves8.1%, while small +chunks regress the output projection. Outputs and changed-input graph replay +retain the normal artifact's bytes. The general candidate fails its8% benefit +gate and is rejected; no runtime policy or native change is added for it. + +## HC independent-row parallelism + +The prior HC launch uses256 threads, with several exchange, row reduction and +copy phases assigning independent rows to eight warps. The16-warp candidate +keeps the same80-CTA resident grid, FP32 reduction order, six down-projection +warps, five up-projection warps, shared-memory lifetime and peer epochs. Extra +warps handle independent row exchanges/reductions and coalesced copies; up +weights and eight-row gate tiles remain guarded to their original warps. +The compiler reports113 registers/thread and43040 static shared bytes for the +512-thread research kernel, without spilling. This is a launch-resource fact, +not a measured occupancy or memory-bandwidth claim. + +Four-rank64-call graph screens preserve bytes with changed inputs. Critical +rank maxima at C4 improve0.064814 to0.056273ms, and at C8 improve0.107693 to +0.096882ms. A384-thread intermediate gives0.059313/0.097945ms. The512-thread +schedule is then folded into the existing native HC batch operation, retaining +its ABI, workspace and runtime M17..40 admission; M9..16 still uses the retained +runtime path while the acceptance-history issue remains under investigation. + +The normal CMake artifact passes44 CPU policy/dispatch cases, the compiled +four-rank payload/state suite and direct changed-input graphs. Direct C1/2/4/6/8 +costs are0.022916 /0.037545 /0.055384 /0.075440 /0.097872ms, versus retained +8-row tiles0.022454 /0.045633 /0.074343 /0.104196 /0.136121ms. C1 is the same +fallback; the direct M10 native result does not widen runtime admission. +The new normal `_C` SHA256 is +`a71013eeb238309d695d7f05977a20ea8681734af041b4671d08d45c16f5c700`. +These local wins require separate whole-model qualification; they do not +establish the requested28ms C8 verifier cost. + +The normal HC16 artifact subsequently completes the capacity-eight model run. +Whole-batch C1/2/4/6/8 verifier means are15.987 /24.517 /32.963 /43.241 / +51.328ms, with187 /228 /172 /161 /178 aligned TP4 rounds and zero incomplete +rounds. Draft means are3.706 /4.460 /5.762 /7.692 /8.738ms. C8 improves from +53.130ms, but C6 regresses from41.812ms. This is not a general MTP speed win; +C2/C4/C8 targets remain unmet. Ordinary common-interval aggregate rates are +138.93 /184.64 /247.54 /274.36 /318.55tok/s, with extra C2/C8 repetitions +184.71 /335.22. All eight natural-EOS arithmetic checks and the C1/C8 +32K retrieval/arithmetic, executable-code and tool/reply gates pass. Actual +worker maps confirm the normal core/indexer binaries without research DSOs. + +## Batched launch amortization + +A separate five-range graph-node trace has163 /121 /96 /70 /61 complete +C1/2/4/6/8 rounds. Traced verifier means20.249 /27.794 /35.067 /48.265 / +55.313ms are perturbed attribution figures, not replacements for the untraced +model timings. Exclusive critical-rank wall categories sum to each verifier. +From C1 to C8, HC grows2.469 to9.749ms, expert projections2.885 to9.473ms, +dense/shared projections2.493 to6.638ms and selected QSA attention1.402 to +3.735ms. Concurrent categories occupy0.998 to9.329ms: that interval represents +overlap, not an overhead that may be added to service sums. Root-only indexer +service must be inspected separately from non-root critical-rank categories. +The trace cannot establish actual occupancy, memory stalls or NVLink bandwidth; +NCU hardware counters are unavailable on this host. + +End-of-capture metadata snapshots show shared-prefix C1/2/4 indexers with +41/43/47 unique physical pages and the same number of distinct leases. Only +1/2/4 tail pages are refreshed. There is no hash-collision evidence to justify +changing the persistent-bank indexing. The remaining per-request score launches +are instead independent M5 computations over shared bank metadata. + +The batched-grid candidate adds an optional normal native +`run_staged_banked_batch` operation to the existing shared-key provider. A +second grid dimension handles uniform five-query request groups, with original +FP16 MMA, FP32 reduction, visibility and paged-overflow semantics. Global live +request IDs replace per-group subtraction. Runtime admission requires the +existing versioned four-tensor bank, C2..8, and exactly five queries per group. +C1, ragged groups, refresh-only staging and older binaries retain their existing +operations. The previous three native names and schemas are unchanged; the +new mutation schema has its own fake registration. No new workspace, provider, +weight/cache format, configuration parameter or default promotion is introduced. + +A normal-flag private graph screen includes page planning, tracked tail writes +and scoring. At32K shared-prefix C1/2/4/6/8, retained/batched means are +0.04474/0.04459,0.06750/0.05509,0.10023/0.07313,0.13306/0.09531 and +0.17483/0.11679ms per indexer. Independent-prefix C8 is2.20451/2.22023ms +(0.7% regression, within the predeclared3% bound); other measured independent +cases improve. All ten cases retain score/visibility bytes with changed cache +writes and live graph metadata. These are private operator screens; normal +artifact and whole-model evidence are recorded separately. + +The MTP Q8_0 draft uses the existing GGUF MoE family. Its retained vector kernel +launches one32-thread CTA per expert output row; the traced C8 down projection +alone launches1,024,000 CTAs. The row-batched operation groups eight independent +output warps in one CTA, retaining activation quantization, dp4a, per-warp +FP32 accumulation, FP16 stores and weighted expert reduction. The shared GGUF +executor keeps the above64-token MMQ path. Native and fused-MoE schemas gain +new optional names while the existing public signatures remain unchanged. +The MTP method binds the operation once after weight loading for SM70 with a +matching native symbol; other hardware or older artifacts retain the previous +executor. This does not alter quantized checkpoint bytes or expert placement. + +A fourteen-case private screen uses actual TP4 projection dimensions,64 +experts and logical M1/2/5/8/10/20/40. All eager and changed-input/routing graph +outputs retain bytes. The down projection improves3.67..4.45x; C8-aligned +M40 drops1.45110 to0.34090ms. Gate/up improves approximately1.08x at M40, +with a1.6% M2 regression within the3% small-case limit. This optimization +addresses draft cost, not the target verifier. Full512-expert weighted output, +normal artifact, capture/AOT, threshold and complete MTP4 gates remain separate. + +The first normal-artifact pass has85 passing cases and one failing M65 +weighted-MoE boundary case. The retained MMQ kernel incorrectly skips every +expert ID above255, leaving output rows uninitialized for this512-expert +model. A smaller32-expert diagnostic is finite across K160/256/2560, ruling +out the initial K-padding hypothesis. The shared MMQ kernel now skips only +negative padding IDs; valid expert IDs still come from the existing checked +route layout. Regression coverage includes expert-count neighbors256/257, +the actual512 count and513, plus full512-expert weighted M64/M65 graphs. +The operation order and MMQ/vector threshold remain unchanged. This fixes +prefill/boundary draft computation; it does not claim a verifier speedup. + +The corrected normal artifacts pass90 GPU/initialization/state cases, including +full512-expert M0/1/2/5/8/10/20/40/64/65 weighted outputs, changed-input and +routing graph replay, row tails1/7/8/9/319/321, fake/AOT dispatch and compiled +fused MoE. Versioned-bank cases now include C4/C6, two position integer types, +noncontiguous Q/table views, ragged request groups and missing-new-op fallback. +The existing approximate dequantized-reference checks also pass. Native byte +comparisons are used where operation order is unchanged; MMQ vs vector +expert-ID coverage uses a predeclared relative norm error bound of0.002. + +The normal512-expert vector-operation screen includes allocation, quantization +and stores. At aligned M40, gate/up drops0.42917 to0.40934ms and down drops +1.45505 to0.34477ms; eager and changed-input graphs retain bytes. The normal +indexer screen reproduces shared-prefix C8 0.17455 to0.11694ms, while independent +C8 is2.19487 to2.19086ms. Whole-model verification remains a separate gate. +Fresh-process dependency inspection loads only normal source-owned artifacts. +The stable-libtorch SHA256 is +`45282c1a7a29f694d5d1832a23c0a0731bf9f3853546582c5c0a14f7bdfec37b`; +the indexer SHA256 is +`ae483cb461ce78a5f37a3c7f52c297f0cdd8d9211eacc4133acba58bf8d39848`. +The HC/core artifact remains the qualified`a71013ee` build. + +## Whole-model indexer-grid and Q8-row result + +Checkpoint`b303bc4c6bf48633541a1a4aad8df9e6372ea802` passes fresh ordinary +workers and a separate fenced diagnostic service. The matched32K shared-prefix +capacity-eight contract, sampling, expert placement and graph sizes are +unchanged. Both actual normal worker maps and startup binding logs confirm the +source-owned core/stable/indexer artifacts. No private research library enters +the service. C1/2/4/6/8 complete-round counts are187/211/192/182/178, with zero +incomplete rounds. + +| C | Verifier before/after, ms | New verifier p99 | Draft before/after, ms | Ordinary aggregate tok/s | +| --- | ---: | ---: | ---: | ---: | +| 1 |15.987 /15.793|18.245|3.706 /3.470|140.98| +| 2 |24.517 /24.216|28.515|4.460 /3.966|190.86 /167.15| +| 4 |32.963 /33.108|43.364|5.762 /4.901|241.54| +| 6 |43.241 /40.769|47.697|7.692 /5.892|302.62| +| 8 |51.328 /49.399|58.014|8.738 /6.710|362.03 /328.00| + +Verifier is the whole TP4 maximum of target forward, sample and state commit; +it is never divided by concurrency. Draft statistics are on the critical +verifier rank. Ordinary rates use a common decode interval; acceptance metrics +cover the entire request batch and cannot be substituted into an exact +per-round cost calculation. The candidate reduces C8 verification by1.929ms +and draft by2.028ms; C4 verification is essentially unchanged/slightly worse. +The requested18/22/28ms C2/C4/C8 verifier targets remain unmet. Natural-EOS +arithmetic, C1/C8 retrieval/arithmetic, executable code and automatic tool/reply +gates all pass. Actual256K and independent-prefix whole-model coverage remain +unqualified. The latest memory sample is after shutdown and is not peak evidence. + +## Subsequent HC and recurrent scheduling decisions + +A private HC16 phase diagnostic keeps the normal arithmetic and uses timestamp +stores only to localize work. At C8, mean CTA phases include approximately +27.47us outbound peer stores,6.10us weight prefetch,14.16us peer wait/combine, +4.11us down MMA/partial stores,18.02us repeated RMS-statistic reduction and +normalization,5.34us down reduction,4.94us LoRA loading and5.24us up/gate-mix. +Grid rendezvous and other waits are separate. These instrumented CTA phases +are perturbed and are not a closed ordinary-worker model latency table. Launch +resources alone do not establish occupancy, stalls or NVLink bandwidth. + +The two-tile HC experiment doubles independent down/up MMA row tiles within +512 threads and uses48320 shared bytes. It retains bytes in four-rank changed +input graphs, but C2/4/6/8 maxima improve only0.7/0.6/0.3/0.5%. C8 is +0.097190 versus0.096711ms, below its predeclared8% benefit gate. The candidate +is rejected before source integration. The hypothesis that serial MMA tiles +explain most HC batch inflation is not supported by the measured phases. + +A separate existing-provider GDN schedule screen covers C1/4/8 and BV/warp +pairs32/4,16/4,32/8,64/8. All twelve cases pass finite eager and10240-token +recurrent-state comparisons with live accepted lengths. Output relative RMS +is bounded to0.001; FP32 state relative RMS0.00003 and max absolute error +0.00002. These operation/dtype-aware limits were declared before the screen. +The best C8 kernel result is0.05904 to0.05503ms (6.8%), below its15% benefit +gate. No runtime scheduling change or default is promoted. + +The next private hypothesis is duplicated HC RMS-statistic reads: every CTA +currently recomputes the same per-row80-way reduction. Any experiment that +publishes a shared result must preserve sum order, reset epochs before the +existing grid rendezvous, use coherent payload accesses with matching GPU +release/acquire, preserve workspace lifetime and retain all fallback paths. +The [PTX memory model](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#memory-synchronization) +provides the ordering contract; normal native state and complete-model gates +are required before accepting a corresponding source change. + +The RMS-owner private screen meets the8% C8 gate while retaining bytes in +four-rank changed-input graphs. C2/4/6/8 critical-rank HC means change +0.039163/0.059407/0.075162/0.097531 to +0.036714/0.052471/0.067874/0.089070ms. The kernel retains512 threads, +113 registers/thread and43040 shared bytes without spills. These are private +operator results, not complete-model gains. + +The candidate is folded into the existing HC batch kernel without changing +its ABI, allocations, peer exchange, fallback policies or M17..40 runtime +admission. One resident CTA owns each row's canonical80-way RMS reduction; +the FP32 sum, rsqrt and FP16 boundaries remain unchanged. It publishes four +inverse values in padded down columns84..85 belonging to CTA79, which have +no downstream down-reduction readers. The last producer resets epoch flags +after its own partial stores and before the existing acq_rel grid barrier. +The owner then uses a GPU release store; consumer leaders acquire the matching +epoch before coherent payload loads and warp broadcast. All80 CTAs remain +resident, so a polling consumer cannot prevent its producer from running. +The first grid barrier carries reset/partial writes; no external initialization +or stale-address assumption is added. Subsequent tile fallbacks may reuse the +workspace; the next wide call resets all of its live row epochs again. + +Normal-artifact admission requires the affected compiled payload/state test, +changed graph inputs and widths, the local cumulative grid counter crossing +uint32 wrap, live inverse/epoch validation and the ordinary MTP4/task gate. +The peer epoch protocol itself is unchanged; wrapping local counters does not +claim peer-epoch wrap qualification. The Q8 row binding also explicitly rejects +non-byte packed weights without changing valid-input computation. + +The normal RMS-owner artifact passes45 interface/admission cases and the +four-rank compiled payload/state suite at M5/8/9/10/16/17/20/24/25/30/32/33/ +39/40/41. The latter includes changed-input replay after local barrier-counter +wrap and validates live inverse values and publication epochs after changing +widths. Direct C1/2/4/6/8 native graph costs are0.022196 /0.035205 /0.050418 / +0.067357 /0.088446ms, with byte-exact retained-tile and changed-input results. +The direct M10 result does not widen runtime admission. Full model gates remain +pending for this artifact. Core SHA256 is +`778595e26ca1c9957ad3f84f8a9d3c1d14a7bbb96e61f33de6353a5307e874c5`; +the Q8 byte-layout guard's stable-libtorch artifact is +`ac770d73a344357f610efd0fd0285e0bedbaedb708e8b76ba46b3d2d1422741f`. +The indexer stays`ae483cb4`. Private operator libraries remain research-only. + +## RMS-owner complete-model result and lossless partial packets + +The RMS-owner checkpoint`a7cf13b080c35abd7294316ee76e882c2db3a3a8` passes its +separate ordinary/task and fenced model gates. Whole C1/2/4/6/8 verifier means +are15.899 /24.508 /31.777 /40.926 /48.937ms; p99 is18.194 /27.965 /41.526 / +51.258 /62.641ms. Draft means are3.485 /3.991 /4.792 /5.910 /6.712ms. +C4 improves1.331ms, C8 improves0.462ms, and C6 is essentially unchanged. +C8 p99 increases from58.014ms, so this is not a universal latency win. Ordinary +rates are140.22 /172.45 /252.33 /309.26 /339.85tok/s; extra C2/C8 repetitions +are157.47 /325.08. Acceptance varies and throughput overlaps previous ranges. +All task and natural-EOS checks pass. Sampled root allocation16078MiB/free67MiB +matches the earlier campaign, without proving transient or256K boundaries. + +The next experiment reduces HC peer bytes losslessly. A sender without a +secondary partial sends two raw FP16 values in one32-bit payload; a sender +with a secondary retains its original FP32 local sum and packet payload. +Per-sender epoch-tagged format words occupy dead LoRA padding324..327, so +mixed optional-secondary presence across ranks remains correct. Every data +word retains the existing64-bit payload+epoch validation; no shared flag for +a multiword payload or new cross-GPU fence assumption is introduced. The +canonical rank sum, reductions and FP16 storage boundaries remain unchanged. +Two scalar half loads support contiguous views aligned only to two bytes. + +The private four-rank C4/C8 graph screen retains bytes with changed inputs. +C8 HC changes0.088271 to0.070946ms (19.6%) without secondary inputs. Mixed +secondary presence regresses0.4%; all-secondary cases regress2.1%, within the +predeclared3% bound. C4 no-secondary improves15.1%. These local figures do +not establish complete-model speed. No workspace or allocator is added. + +Normal integration adds`sm70_hcx_batch_packed_out`; the prior batch and small +native names/schemas retain their packet formats. The existing HC kernel is +templated on the internal packet schedule rather than duplicated as a new +provider. Initialization binds the new optional symbol once; older native +artifacts retain the existing batch operation. Fake registration and missing +symbol tests are separate from actual GPU execution. All TP workers must use +the matching normal artifact, as for other collective kernels. + +C2 native-batch admission is requalified after the corrected512-expert MMQ +prefill boundary. For an already initialized capacity-forty HC owner, separate +projections may now use native M9..40; M1..8 and fused producers retain tiles. +Startup configurations with maximum at most16 keep the previous smaller owner +and tiled fallback, so this does not add an unmeasured memory requirement to +those modes. The earlier C2 acceptance regression is retained as negative +evidence; full-model acceptance/task and verifier gates are required before +claiming this new admission is qualified. Numerical equivalence alone does +not establish emitted-token throughput. + +The matching normal core SHA256 is +`4fab324fc4082a3206523eafb7cb29f565c39cef4c93c6b0eb5165094d7b639d`. +Fresh-process inspection confirms the normal core/stable/indexer libraries +without private research DSOs or preload. Interface, fallback and focused +M8-neighbor cases pass; the four-rank compiled-payload/state test and raw fake +registration pass. Raw comparisons include M9/10/16/17/25/40 with unaligned +contiguous FP16 views, absent/mixed secondary inputs, changed graph inputs and +FP32 local sums beyond65504 that cancel before the final FP16 boundary. + +Normal paired graph costs at C4/C8 without secondary inputs are +0.055074/0.091793 to0.046548/0.071998ms. Mixed-input changes are-0.1%/+0.3%, +and all-secondary regressions2.2%/1.7%, within the declared bounds. The normal +new operator and retained operator remain callable on the same owned buffers; +no task-cache sidecar is needed. Stable-libtorch`ac770d73` and indexer`ae483cb4` +are unchanged. Full changed-source hooks pass. Whole-model qualification of +the packet schedule and renewed C2 admission is pending. diff --git a/examples/serve_flashnext_gguf_256k.sh b/examples/serve_flashnext_gguf_256k.sh new file mode 100755 index 0000000000..f1a8b8161d --- /dev/null +++ b/examples/serve_flashnext_gguf_256k.sh @@ -0,0 +1,100 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Qualified Flash-Next GGUF TP4 settings: 256K total tokens, host KV and disk PLE. +set -euo pipefail +SOURCE_DIR=${SOURCE_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)} +PYTHON_BIN=${PYTHON_BIN:-$SOURCE_DIR/.venv/bin/python} +: "${MODEL_FILE:?Set MODEL_FILE to the first IQ3_S GGUF shard}" +: "${MTP_DIR:?Set MTP_DIR to the matching stock MTP safetensors directory}" +: "${CACHE_DIR:?Set CACHE_DIR to a dedicated runtime cache}" +: "${VLLM_API_KEY:?Set VLLM_API_KEY for the public API}" +PORT=${PORT:-8000} +GPU_IDS=${GPU_IDS:-4,5,6,7} +CPU_SET=${CPU_SET:-24-47,72-95} +EXPERT_PROFILE=${EXPERT_PROFILE:-} +QSA_HOST_INDEXER_HISTORY=${QSA_HOST_INDEXER_HISTORY:-false} +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.8} +export PATH="$CUDA_HOME/bin:$(dirname "$PYTHON_BIN"):$PATH" +export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES="$GPU_IDS" +export PYTHONPATH="$SOURCE_DIR" +export TORCH_CUDA_ARCH_LIST=7.0 +export PYTORCH_ALLOC_CONF=expandable_segments:True OMP_NUM_THREADS=1 MALLOC_ARENA_MAX=2 +export VLLM_NO_USAGE_STATS=1 DO_NOT_TRACK=1 +export VLLM_EXPERT_CACHE_STACKS=1 VLLM_BALANCED_PREFILL_CHUNKS=1 +export VLLM_SM70_MTP_PROFILE=0 VLLM_NVTX_SCOPES_FOR_PROFILING=0 +export MAX_JOBS=2 TORCHINDUCTOR_COMPILE_THREADS=1 +export TOKENIZERS_PARALLELISM=false +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 +export VLLM_CACHE_ROOT="$CACHE_DIR/vllm" TRITON_CACHE_DIR="$CACHE_DIR/triton" +export TORCHINDUCTOR_CACHE_DIR="$CACHE_DIR/inductor" +export TORCH_EXTENSIONS_DIR="$CACHE_DIR/torch_extensions" +export VLLM_SM70_GEMM_LUT_PATH="$CACHE_DIR/gemm-lut-{device}.bin" +unset LD_PRELOAD LD_AUDIT ONECAT_VLLM_EXTENSION_DIR +mkdir -p "$CACHE_DIR" +KC=$("$PYTHON_BIN" - "$EXPERT_PROFILE" "$QSA_HOST_INDEXER_HISTORY" <<'PY' +import json +import sys +if sys.argv[2].lower() not in ("true", "false"): + raise ValueError("QSA_HOST_INDEXER_HISTORY must be true or false") +config = { + "sm70_gguf": { + "expert_storage": "original", "expert_arena": True, + "embedding_storage": "original", "dense_storage": "canonical", + "small_m_dp4a": True, "q8_expert_intermediate": True, + "small_m_hmma": True, "lut4_expert_dp4a": True, + "device_transcode": True, "dequant_workspace_bytes": 33554432, + "expert_cache_fraction": 0.8, "expert_cache_host_slack": 0.05, + "expert_cache_prefill": "turbomind", "expert_cache_prefill_tokens": 4096, + }, + "sm70_turbomind_workspace_bytes": 8388608, + "sm70_router_weight_storage": "row_major", "sm70_mtp_q8_0_experts": True, + "hc_weight_storage": "sharded", "hc_ll_shard": True, + "hc_ll_optimized_loads": False, "sm70_hcx": True, + "sm70_hcx_output_projection": False, "sm70_hcx_local_schedule": True, + "sm70_qsa_shared_key": True, "sm70_qsa_device_history": True, + "sm70_qsa_host_native": True, "qsa_host_kv": True, + "qsa_host_indexer_history": sys.argv[2].lower() == "true", + "qsa_host_kv_dtype": "fp8_e4m3", + "qsa_host_kv_draft_dtype": "float16", "qsa_host_kv_device_reference": False, + "qsa_host_kv_hot_tokens": 32768, "qsa_host_kv_hot_dtype": "fp8_e4m3", + "qsa_host_kv_draft_hot_tokens": 16384, "qsa_host_kv_staging_rows": 8, + "qsa_host_kv_prefill_stage_tokens": 131072, + "sm70_sparse": {"qsa_score_tile_mb": 32}, "qsa_host_kv_state_blocks": 29, + "qsa_auto_e4m3": False, "ple_row_cache_mib": 512, + "ple_input_prepare": True, "ple_pinned_decode": False, + "ple_fast_lane": True, "ple_draft_prefetch": True, + "sm70_greedy_verify": True, "sm70_sparse_target_sampling": True, + "sm70_sparse_target_reference_check": False, "sm70_draft_single_graph": True, + "sm70_fused_side_projections": True, "sm70_top1x": True, + "sm70_draft_hot_vocab": 0, +} +if sys.argv[1]: + config["sm70_gguf"]["expert_cache_profile"] = sys.argv[1] +print(json.dumps(config)) +PY +) +SPEC=$("$PYTHON_BIN" - "$MTP_DIR" <<'PY' +import json +import sys +print(json.dumps({ + "method": "mtp", "model": sys.argv[1], "num_speculative_tokens": 4, + "draft_load_config": {"load_format": "safetensors"}, + "draft_sample_method": "greedy", +})) +PY +) +cd "$SOURCE_DIR" +exec taskset -c "$CPU_SET" "$PYTHON_BIN" -m vllm.entrypoints.cli.main serve "$MODEL_FILE" \ + --served-model-name flash-next-abliterated flash-next qwen3.8-flash-next \ + --host 127.0.0.1 --port "$PORT" --quantization gguf --tensor-parallel-size 4 \ + --dtype half --kv-cache-dtype float16 --mamba-ssm-cache-dtype float32 \ + --max-model-len 262144 --max-num-seqs 1 --max-num-batched-tokens 8192 \ + --gpu-memory-utilization 0.95 --enable-prefix-caching \ + --mamba-cache-mode align --no-mamba-tail-checkpoint --language-model-only \ + --enable-auto-tool-choice --tool-call-parser qwen3_coder --reasoning-parser qwen3 \ + --default-chat-template-kwargs '{"enable_thinking":true}' \ + --override-generation-config '{"temperature":0.6,"top_p":0.95,"top_k":20}' \ + --compilation-config '{"mode":3,"cudagraph_mode":"FULL","cudagraph_capture_sizes":[5],"inductor_compile_config":{"benchmark_combo_kernel":false}}' \ + --kernel-config "$KC" --speculative-config "$SPEC" + diff --git a/tests/config/test_flashnext_mtp_batch_policy.py b/tests/config/test_flashnext_mtp_batch_policy.py new file mode 100644 index 0000000000..7e2d869104 --- /dev/null +++ b/tests/config/test_flashnext_mtp_batch_policy.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch limits are typed engine policy and survive worker serialization.""" + +import copy +import multiprocessing + +import pytest + +from vllm.config.kernel import KernelConfig + + +def _send_worker_policy(connection, policy): + connection.send(policy) + connection.close() + + +def test_mtp_batch_policy_is_owned_hashed_and_serializable(): + legacy = KernelConfig() + tiled = KernelConfig( + sm70_hcx=True, + sm70_hcx_batch_max_m=40, + sm70_hcx_batch_fusion=True, + qsa_host_kv_state_compact=True, + sm70_qsa_shared_key_batch=True, + sm70_gguf={ + "small_m_dp4a": True, + "small_m_dp4a_max_m": 40, + "small_m_hmma_batch": True, + "grouped_gate_up": True, + "grouped_gate_up_max_m": 40, + "grouped_down": True, + }, + qsa_host_kv=True, + qsa_host_indexer_history=True, + sm70_qsa_indexer_single_rank=True, + sm70_qsa_indexer_stage_keys=True, + sm70_qsa_indexer_persistent_keys=True, + ) + tiled.sm70_gguf.active = True + context = multiprocessing.get_context("spawn") + receiver, sender = context.Pipe(duplex=False) + worker = context.Process(target=_send_worker_policy, args=(sender, tiled)) + worker.start() + sender.close() + try: + assert receiver.poll(60), "Worker policy transfer timed out" + transferred = receiver.recv() + finally: + receiver.close() + worker.join(10) + if worker.is_alive(): + worker.terminate() + worker.join() + assert worker.exitcode == 0 + assert legacy.sm70_hcx_batch_max_m == 8 + assert legacy.sm70_gguf.small_m_dp4a_max_m == 20 + assert legacy.sm70_gguf.grouped_gate_up_max_m == 20 + assert not legacy.qsa_host_kv_state_compact + assert transferred == tiled + assert transferred.compute_hash() == tiled.compute_hash() + for field in ( + "sm70_hcx_batch_max_m", + "sm70_hcx_batch_fusion", + "qsa_host_kv_state_compact", + "sm70_qsa_shared_key_batch", + "sm70_qsa_indexer_single_rank", + "sm70_qsa_indexer_stage_keys", + "sm70_qsa_indexer_persistent_keys", + ): + changed = copy.deepcopy(tiled) + setattr(changed, field, getattr(legacy, field)) + assert changed.compute_hash() != tiled.compute_hash() + transferred.sm70_gguf.small_m_dp4a_max_m = 20 + assert transferred.compute_hash() != tiled.compute_hash() + + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.grouped_gate_up_max_m = 20 + assert transferred.compute_hash() != tiled.compute_hash() + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.grouped_down = False + assert transferred.compute_hash() != tiled.compute_hash() + + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.grouped_gate_up = False + assert transferred.compute_hash() != tiled.compute_hash() + assert tiled.sm70_gguf.small_m_dp4a_max_m == 40 + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.small_m_hmma_batch = False + assert transferred.compute_hash() != tiled.compute_hash() + + +def test_disabled_parents_keep_inactive_batch_options_out_of_hash(): + a = KernelConfig(sm70_hcx=False, sm70_qsa_shared_key=False) + b = KernelConfig( + sm70_hcx=False, + sm70_qsa_shared_key=False, + sm70_hcx_batch_max_m=40, + sm70_hcx_batch_fusion=True, + sm70_qsa_shared_key_batch=True, + sm70_qsa_indexer_stage_keys=True, + sm70_qsa_indexer_persistent_keys=True, + ) + assert a.compute_hash() == b.compute_hash() + + +@pytest.mark.parametrize( + "host, history", [(False, False), (False, True), (True, False)] +) +def test_single_rank_indexer_is_inactive_without_host_history(host, history): + a = KernelConfig(qsa_host_kv=host, qsa_host_indexer_history=history) + b = KernelConfig( + qsa_host_kv=host, + qsa_host_indexer_history=history, + sm70_qsa_indexer_single_rank=True, + sm70_qsa_indexer_stage_keys=True, + sm70_qsa_indexer_persistent_keys=True, + ) + assert a.compute_hash() == b.compute_hash() + + +@pytest.mark.parametrize("limit", [0, 41]) +def test_mtp_batch_policy_rejects_unsupported_limits(limit): + with pytest.raises(ValueError): + KernelConfig(sm70_hcx_batch_max_m=limit) + with pytest.raises(ValueError): + KernelConfig(sm70_gguf={"small_m_dp4a_max_m": limit}) + + +@pytest.mark.parametrize("limit", [19, 21, 39, 41]) +def test_grouped_gate_policy_rejects_unqualified_fragments(limit): + with pytest.raises(ValueError): + KernelConfig(sm70_gguf={"grouped_gate_up_max_m": limit}) diff --git a/tests/kernels/quantization/test_gguf_compact_storage.py b/tests/kernels/quantization/test_gguf_compact_storage.py index 650951d4c1..736b8a73ff 100644 --- a/tests/kernels/quantization/test_gguf_compact_storage.py +++ b/tests/kernels/quantization/test_gguf_compact_storage.py @@ -100,6 +100,8 @@ def test_original_experts_keep_fast_route_for_graph_padding(monkeypatch, m): method = object.__new__(GGUFNativeMoEMethod) method.dp4a_admitted = True + method.expert_cache = None + method.expert_trace_tokens = 0 method.weight_types = {"w1": 21, "w3": 21, "w2": 42} method.input_padding = {"w2": (32, 192)} method.q8_intermediate = True diff --git a/tests/kernels/quantization/test_gguf_dense_segments.py b/tests/kernels/quantization/test_gguf_dense_segments.py index b7307fb747..5ecaf4a929 100644 --- a/tests/kernels/quantization/test_gguf_dense_segments.py +++ b/tests/kernels/quantization/test_gguf_dense_segments.py @@ -114,6 +114,92 @@ def test_reject_m33_without_launch(): ) +@pytest.mark.parametrize("source", [8, 12, 13, 14, 20, 23]) +@pytest.mark.parametrize( + "m,k,n,split,warps", + [ + (9, 2560, 512, 1, 4), + (10, 2560, 4096, 1, 4), + (17, 2560, 3584, 1, 8), + (20, 2560, 3328, 1, 4), + (25, 1536, 2560, 1, 8), + (30, 2560, 160, 5, 8), + (33, 2560, 4096, 1, 4), + (40, 2560, 3328, 1, 4), + ], +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_batch_weight_reuse_official_and_changed_graph(source, m, k, n, split, warps): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + import vllm._C # noqa: F401 + + raw = raw_weights(source, n=n, k=k) + fmt, q, scale, minimum, group = decode(raw, source) + payload = [torch.from_numpy(p).cuda() for p in pack(fmt, q, scale, minimum, group)] + x = torch.randn(m, k, device="cuda", dtype=torch.float16) * 0.125 + output = torch.empty(m, n + 5, device="cuda", dtype=torch.float16)[:, 3 : n + 3] + tiles = (n + 31) // 32 + ws = torch.empty( + tiles * ((m + 7) // 8) * split * 256, device="cuda", dtype=torch.float32 + ) + counter = torch.zeros(tiles * ((m + 7) // 8), device="cuda", dtype=torch.int32) + official = torch.from_numpy(reconstruct(fmt, q, scale, minimum, group)).cuda() + + def run(): + torch.ops._C.gguf_dense_segments_batch_sm70_out( + x, + [payload[0]], + [payload[1]], + [payload[2]], + [output], + [fmt], + [n], + k, + split, + warps, + ws, + counter, + None, + ) + + run() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for shift in (0.0, 0.125): + x.normal_(mean=shift, std=0.125) + graph.replay() + reference = x.float() @ official.float().T + torch.accelerator.synchronize() + assert bool(torch.isfinite(output).all()) + assert float((output.float() - reference).norm() / reference.norm()) < 0.003 + assert not bool(counter.any()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_batch_rejects_m41_before_launch(): + import vllm._C # noqa: F401 + + with pytest.raises(RuntimeError, match="1..40"): + torch.ops._C.gguf_dense_segments_batch_sm70_out( + torch.empty(41, 256, device="cuda", dtype=torch.float16), + [], + [], + [], + [], + [], + [], + 256, + 1, + 4, + torch.empty(1, device="cuda"), + torch.empty(1, device="cuda", dtype=torch.int32), + None, + ) + + @pytest.mark.parametrize("source", [8, 12, 13, 14, 20, 23]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_transient_restore_matches_turbomind_storage(source): diff --git a/tests/kernels/quantization/test_gguf_grouped_down.py b/tests/kernels/quantization/test_gguf_grouped_down.py new file mode 100644 index 0000000000..2d7fed6116 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_grouped_down.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Grouped original down retains TP block boundaries and route-sum order.""" + +import numpy as np +import pytest +import torch +import vllm._C # noqa: F401 + +from tests.kernels.quantization.test_gguf_dp4a import quantize_reference +from vllm.model_executor.layers.quantization.gguf_grouped_gate import grouped_down +from vllm.model_executor.layers.quantization.gguf_native import packed_tp_span +from vllm.transformers_utils.gguf_tensor_reader import dequantize, quant_size + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("source", [20, 42]) +@pytest.mark.parametrize("rank", [0, 1, 2, 3]) +@pytest.mark.parametrize("m", [25, 30, 40]) +def test_grouped_down_changed_graph_exact_and_official_oracle(source, rank, m): + experts, n, full_k, k, top_k = 512, 32, 640, 160, 10 + block, size = quant_size(source) + rng = np.random.default_rng(20261015 + source) + raw = rng.integers(0, 256, (experts, n, full_k // block, size), dtype=np.uint8) + raw[..., :2] = np.array([0.002], dtype="