CUDA/HIP: partial-top-k + hipCUB large-N for argsort/top-k on ROCm#15
Open
DeanoC wants to merge 17 commits into
Open
CUDA/HIP: partial-top-k + hipCUB large-N for argsort/top-k on ROCm#15DeanoC wants to merge 17 commits into
DeanoC wants to merge 17 commits into
Conversation
This was referenced Jul 9, 2026
b-albar
approved these changes
Jul 10, 2026
…Org#501) * vendor: grouped-expert MUL_MAT_ID for small MoE verify batches Speculative-verify batches (2-16 tokens) route consecutive draft tokens to heavily overlapping expert sets, but the multi-token MoE MMVQ kernel reads each (token, slot) pair's expert weights independently. Measured on a 256-expert top-8 MoE, an 8-row verify batch reads 64 expert matrices per layer where only 39.4 distinct ones are needed (1.62x duplication); at 15 rows it is 120 vs 58.9 (2.04x). This adds an opt-in path (DFLASH_MMID_GROUPED=1) that rank-sorts the (token, slot) pairs by routed expert in a tiny capture-safe prep kernel, then runs a variant of mul_mat_vec_q_moe whose warp w processes sorted pair blockIdx.y*8+w instead of (slot, token). Same-expert pairs are adjacent after sorting, so the warps of a block share an expert and duplicate weight reads are served from L1/L2 instead of DRAM. Launch shape, per-warp structure, vec_dot sequence and warp reduction are identical to mul_mat_vec_q_moe, so outputs are bit-exact and CUDA-graph capture stays supported. DFLASH_MMID_GROUPED_TYPES bitmask selects quant types: 1 = Q4_K (default), 2 = Q6_K (also lifts the Q6_K MUL_MAT_ID MMVQ ceiling to 16, keeping CUDA graphs enabled for those batches), 4 = Q4_0/Q8_0/Q5_K. Model-agnostic for any MoE with n_expert_used <= 16. Measured on RTX 3090 (sm_86), Laguna XS 2.1 Q4_K_M target + DFlash drafter, greedy HumanEval prompts, outputs byte-identical: verify width 8: 212.0 -> 219.9 tok/s (+3.7%) verify width 15: ~157 -> 173.3 tok/s (~+10%) Gains grow with verify width (more expert overlap to dedup), which is what makes wide/hedged verify configurations cheaper. * vendor+server: opt-in per-token adaptive expert count for MoE verify batches Speculative-verify tokens differ in difficulty: for most tokens the router mass concentrates in a few experts, and dropping the low-mass tail rarely changes the verified argmax. This adds DFLASH_ADAPTIVE_K_TAU=<0..1>: inside the grouped MUL_MAT_ID prep kernel, each token keeps its leading experts until their cumulative combine weight reaches tau; the remaining ids are sentineled to -1 (skipped by the grouped kernel, exact zero contribution) and the kept weights are renormalized in place. Requires DFLASH_MMID_GROUPED=1 with DFLASH_MMID_GROUPED_TYPES=7 so every routed quant type takes the grouped path. The graph-builder side tags the router ids tensor with the combine-weights tensor via ids->extra (server/src/common/mmid_adaptive_k.h). Wired for laguna (DFlash capture layers 1,13,25,33,39 kept dense by default), qwen35moe and gemma4 (dense layers via DFLASH_ADAPTIVE_K_DENSE). Prefill and single-token decode are untouched. Not bit-exact (near-lossless): tau trades a small per-token risk for speed. Measured on RTX 3090, Laguna XS 2.1 Q4_K_M + DFlash drafter, HumanEval-10 greedy, tau=0.80: 212.0 -> 224.5 tok/s (+6% combined with the grouped kernel), expected_pass 10/10 at every tau tested; tau=0.80 best. tau needs per-model validation; qwen35moe/gemma4 are mechanism-wired but not yet swept. * vendor+server: drafter-confidence adaptive verify width for MoE spec decode Every verify row on an MoE target routes to its own top-k experts, so rows cost real weight bandwidth (~1ms/row for a Q4 target on a 3090). A fixed verify width pays that on every step; this trims the batch per step using the drafter's own confidence (keep slot j while the product of top-1 probs through j stays above DFLASH_ADAPTIVE_WIDTH_THETA). Output-exactness is preserved: only the number of speculated slots changes. - ggml-cuda: ggml_backend_cuda_topk_rows, a small top-k (k<=8) + softmax probability reduction over draft-head logits rows, ~KB readback instead of the full vocab logits. Also used to accelerate project_hidden_to_topk (draft-tree candidates) at unit temperature. - common: adaptive_verify_width.h helper + DFlashTarget hook returning per-slot candidate probs alongside draft tokens; domino fused head exposes its corrected logits for the same extraction. - laguna: reference integration (projection extraction, per-width verify graph slots so every width's CUDA graph captures once and replays, trim before verify_batch). Laguna XS 2.1 Q4_K_M + official DFlash drafter, 3090, HE-10 512-tok: theta=0.20 min=4 --verify-width 8 -> 236.0 tok/s vs 226.4 fixed w6 control (+4.2%), commit/step 3.58->3.79, avg rows 5.35, 10/10 pass. Off by default. * server: make validated MoE verify optimizations the default, single flag for the rest Collapse the env sprawl from the previous three commits into a clean surface: - Grouped MUL_MAT_ID: bit-exact, so ON by default on CUDA (HIP stays opt-in, unvalidated). DFLASH_MMID_GROUPED=0 is the kill switch; the type mask defaults to all validated types and is a debug override. - Adaptive verify width: output-exactness-preserving, so ON by default (theta 0.20, min 4). Greedy chains now default to a base width of 8 rows trimmed per step by drafter confidence; the legacy accept-EWMA AUTO remains the fallback for theta=0, sampled verify and --ddtree. - Adaptive expert count (the only output-changing knob): promoted from env to --adaptive-experts [tau] (default 0.80). Explicit env still wins for sweeps. Net: a default server run gets the exact optimizations with zero configuration; one documented flag opts into the near-lossless expert gating. * fix(draft): load and apply Laguna XS 2.1 drafter tensors (attn gates, aux hidden norms, context-KV norms) The official Laguna XS 2.1 DFlash drafter carries per-layer attention gates (blk.<i>.attn_gate.weight), per-capture-layer aux hidden norms (dflash.aux_hidden_norm.<n>.weight) and a context-KV layer input norm flag. The draft loader on main silently ignored all three, so the drafter ran without its feature norms: the fc projection overflowed in f16 and every draft hidden state came out NaN, killing spec decode on any main build (laguna_verify_batch: embed failed on token id -1 from an all-NaN argmax). Port the loader + draft-graph support so main can run the official drafter. * perf(spec): device-resident domino head, persistent draft graph, drafter requant tool - fused Domino head reads the draft graph's device hidden in place (hidden_dev param); the 64KB D2H hidden readback is now lazy and only paid by fallback heads / ddtree / sampled-verify paths - build_draft_step skips the per-step rebuild while ctx stays inside the same 64-aligned pad bucket (copy-mode builds only); laguna passes copy-mode by default, kill with DFLASH_DRAFT_PERSIST=0 - scripts/requant_dflash_draft.py: GGUF->GGUF drafter requant (q8/q4), metadata verbatim; q4 drafter measured +3% end-to-end on Laguna XS 2.1 (3090), acceptance unchanged - drafter is latency-bound, not bandwidth-bound - step-prof: commit/build laps close the loop-top blind spot * feat(spec): drafter context-KV ring cache (target-agnostic) DFlash drafters re-encoded their whole feature window (up to draft_ctx_max tokens of target hidden states) every step: ~10ms/step once the window fills, on every model family, previously hidden behind slower targets. common/dflash_draft_kv.{h,cpp}: per-layer head-major F16 ring caches (slot = pos % cap) with absolute-position RoPE, so entries stay valid as the window slides. One fixed-topology step graph (fold-in append of newly committed rows via set_rows with a trash slot for pads, then the noise forward flash-attending over the ring) built once, CUDA-graph replayed forever; bulk append after prefill. All inputs live in a dedicated buffer so gallocr never recycles them. draft/draft_graph.cpp: build_draft_kv_append/build_draft_kv_step next to the legacy one-shot builder, sharing draft_fuse_features. The per-position k_norm and RoPE are separable, so append/step splitting matches the legacy concat math exactly (f16 storage is the only numeric difference). Any backend adopts it with ~20 lines: draft_kv_init once, then per step draft_kv_begin_step + upload noise embed + compute. Measured (laguna XS 2.1, RTX 3090): draft step 9.9 -> 2.4ms, flat at any context; acceptance parity. * feat(laguna): flat long-context serving under KVFlash Laguna XS 2.1 Q4_K_M + q4 drafter on RTX 3090, cold-certified: decode 154.6 tok/s @30k and 152.3 @256k (8K pool, was 84.8/22.1), prefill 240K tokens in 67.3s (was 26.4 min), acceptance parity, ~1GB VRAM freed. - SWA ring caches [TAG_SWA_RING]: 30/40 layers attend a 512-token window but flash-attended the whole pooled span with 97% masked out. Window layers now keep a position-indexed ring (slot = pos % ring, ring sized from chunk+window) and bypass the pager entirely; the pager only pages full-attention layers. Verify 19.4 -> 13.7ms, prefill -27%. Snapshots are skipped when rings are active (they hold only the trailing window); kv_cap probes pick a full-attention layer. Kill: DFLASH_LAGUNA_SWA_RING=0. - Drafter ring-cache wiring behind DFLASH_DRAFT_KV (default on, =0 legacy; rebuilt on drafter-variant switch). - Pooled prefill batches multiple pager chunks per target step and the feature capture appends via set_rows with ring-row indices as input data, so prefill graphs CUDA-replay; --chunk 1024 feeds the MoE expert GEMMs better (-19% prefill wall at 240K). - Env-gated profilers: DFLASH_LAGUNA_STEP_PROF (commit/build/dwait laps), DFLASH_LAGUNA_VERIFY_PROF and DFLASH_LAGUNA_PREFILL_PROF (sub-phase laps), GGML_CUDA_GRAPH_STATS_EVERY + first-mismatch logging in ggml-cuda. Measured: the entire KVFlash host machinery costs 0.08ms/step; verify is 99.5% GPU kernel time. * feat(qwen35): adopt the drafter context-KV ring cache Wires the shared common/dflash_draft_kv module into the qwen35 backend's local-draft path (same pattern as laguna: lazy init, built_for guard on drafter switch, DFLASH_DRAFT_KV=0 kill switch, legacy path preserved). Remote-draft (IPC) requests keep the legacy path. Also adds distinct error messages to the bulk-append failure paths in dflash_draft_kv.cpp; a mismatched drafter now fails gracefully with the offending dims where the legacy path hit a GGML_ASSERT OOB write. A/B on Qwen3.6-27B Q4_K_M + dflash-draft-3.6-q4_k_m (RTX 3090): short 83.7 -> 96.3 tok/s (accept 38.0 -> 41.2%), 8K real code 26.8 -> 31.6 tok/s (19.3 -> 18.7%), outputs coherent both legs. * fix(spec): address review findings on the drafter ring cache - Reset the drafter ring at the start of every spec-decode request in both backends: the state persists across requests, so a new conversation could draft against the previous request's cached rows for positions below next_pos. The first begin_step re-appends the live window from the feature mirror (~10ms once per request). Validated: identical speed, acceptance and output across back-to-back requests. - StepGraph records built_view; the copy-mode persistent fast path refuses graphs whose topology reads features through a mirror ring view (reachable with DFLASH_DRAFT_PERSIST=0). - DraftKvState is now non-copyable: it owns raw ggml contexts, buffers and graphs, so a copy would double-free. - Document that the drafter SWA window is anchored at committed for all noise rows on purpose: it replicates the legacy one-shot graph's single windowed view, which is the trained drafter's validated attention pattern. * fix(review): harden the mmid/adaptive-K/ring surfaces for merge - Adaptive-K gates a scratch COPY of the router ids: the -1 drop sentinels never reach consumers on the non-grouped path (a sibling weight with a non-grouped quant type would cast them to uint32_t and read OOB). The shared zeroed gate weights keep every consumer bit-correct, and the already-gated marker is now weights-based since ids copies are per-op. - Grouped MUL_MAT_ID falls back to the legacy kernel above MMID_GROUPED_MAX_PAIRS instead of aborting on a hard assert. - q8 memoization re-enabled for MUL_MAT_ID (quantization is ids-independent and the memo key never included ids). - SWA ring guards: build_laguna_graph asserts ring => kv_idx_swa; laguna_layer_step refuses ring caches (unwired path fails loudly, not OOB). - Adaptive-K extras pooled per tensor address (was one leak per graph build); tau/DENSE/theta/min env knobs and --adaptive-experts get strict validated parsing; il<0 families warn once that dense exclusions cannot apply yet (layer-index threading tracked as follow-up). - Draft GGUF loader frees the metadata arena on all early validation failures; requant tool refuses src==dst (in-place corruption). - GGML_CUDA_GRAPH_STATS_EVERY clamped positive (division by zero); topk-rows sync contract + output buffer sizes documented. Validated: laguna ring A/B per-step parity (17.2/22.8 ms per step vs 17.3/23.1 pre-fix), opt-in grouped+adaptive-K leg 230 tok/s short / 126.7 @8K, coherent output, acceptance healthy. * refactor(env): consolidate profilers into DFLASH_PROF, document the env surface The server had accumulated 153 distinct environment variables. Policy going forward (docs/ENVIRONMENT.md): features ship as CLI flags or defaults; env vars are reserved for burn-in kill switches (deleted after soak) and debug instrumentation. - DFLASH_PROF=step,verify,prefill replaces DFLASH_LAGUNA_{STEP,VERIFY, PREFILL}_PROF (all three were new this branch; no users to migrate). - docs/ENVIRONMENT.md documents the load-bearing variables (default, purpose, lifecycle) and carries the full generated inventory; the promotion/deletion audit of the legacy surface is tracked as follow-up. * polish(review): log-spam guard, honest comments, doc rendering Self-review pass over the review-fix batch: - malformed DFLASH_ADAPTIVE_K_DENSE entries warn once, not per graph build - the adaptive-K extras pool documents its single-threaded-builder assumption (same contract as the thread_local builder arenas) and states the boundedness argument honestly - two profiler comments still referenced the pre-consolidation env names - ENVIRONMENT.md: a raw pipe inside a table cell broke markdown rendering Validated on a full clean recompile: DFLASH_PROF=step emits the step profile (the rename's runtime behavior, previously masked by a stale object: rsync -a preserves source mtimes older than build objects, so make skipped the recompile - box-tree-only issue, the pushed branch was always verified via the fresh-clone build), opt-in grouped+adaptive leg byte-identical behavior. * fix(deploy): no tq3_0 or finite fa-window in any default path Found by running the hermes harness against laguna: harness defaults forced tq3_0 KV (garbles laguna output) and fa-window 2048 (drops the system prompt, breaks tool calls). - server: the max_ctx>6144 auto-tq3_0 policy is removed; KV types come from each family's default (laguna q8_0, base q4_0), tq3_0 only via an explicit --cache-type flag - laguna: explicit tq3_0/q4_0 KV now prints a loud garbling warning - harness/clients/common.sh: no KV override exported unless the user sets one; FA_WINDOW defaults to 0 (full attention) Validated: hermes harness run against laguna shows family-default KV and coherent output (the previous run produced gibberish). * feat(laguna): tool calling end to end Laguna emits poolside's tool dialect - <tool_call>NAME with <arg_key>K</arg_key><arg_value>V</arg_value> pairs - where the wrapper tags are special tokens that detokenization strips. Nothing in the pipeline knew the format, so agent harnesses saw plain text and dead-ended after one turn. - chat_template: the laguna system block renders the tools declaration per the GGUF-embedded template (### Tools, schemas as raw JSON inside <available_tools>, calling instruction with the arg_key/arg_value example, thinking + non-thinking variants). Fixes the model inventing tool/arg names it could not see. - tool_parser: pattern 8 parses both the wrapped and the stripped form (name walked back from the <arg_key> anchor; values typed through the declared schema like the other seven dialects). - sse_emitter: <arg_key> added to the tool triggers with a name rewind, and the safe-prefix flush holds back a trailing identifier run (<=64 chars) when tools are declared so a streaming split cannot separate the tool name from its first tag. Validated on RTX 3090: the probe returns finish_reason=tool_calls with {"name":"get_weather","arguments":{"city":"Paris"}} and clean content; parser verified standalone against the exact model output. * refactor(tools): hoist the duplicated trim lambda into trim_ws Pure refactor from the self-review pass; parser behavior verified identical on the standalone test (get_weather/{city:Paris}, clean text). * fix(harness): purge finite fa-window and tq3_0 defaults everywhere Follow-up sweep after common.sh: run_openclaw overrode FA_WINDOW back to 2048, the lucebox-vs-llamacpp benchmark defaulted 2048 + tq3_0 KV, and client_test_runner's launch profiles hardcoded both. All now default to full attention and family-default (or q8_0) KV; tq3_0 remains only in the explicit cache-type comparison matrix, which tests it on purpose. * fix(tools): laguna dialect review findings - Wrapped Laguna bodies whose <arg_value> holds JSON (the template serializes non-string args via tojson) were rejected by the '{' guard meant for Qwen JSON bodies; bodies containing <arg_key> pairs are now accepted. - convert_param_value assumed a string "type"; JSON Schema array types (["integer","null"]) threw out of the parser. First non-null string entry is used instead - hardens all eight dialects. - Zero-argument calls in the stripped form (bare declared tool name at end of output, wrapper tokens removed by detokenization) are recognized at finish and re-wrapped for the parser. Fourth reviewer suggestion (drop the tool-call instruction from the prompt) dismissed: the instruction is a verbatim copy of the model's own GGUF-embedded chat template; removing it would deviate from the trained format. Validated standalone (JSON-object args, array-typed schema, zero-arg) and live on RTX 3090: get_weather/{city:Paris} regression clean, zero-arg list_files/{} returns finish_reason=tool_calls. * fix(server): status report no longer claims the removed tq3_0 auto-default The /status model card still computed 'tq3_0 above 6144 ctx' for the KV type when no explicit --cache-type was passed, describing the auto policy that 5469578 deleted. Report 'family default' like the startup banner. * docs(laguna): drop stale reference to the removed tq3 auto policy --------- Co-authored-by: mrciffa <davide@cifarelli.tech>
Updated performance metrics for Laguna-XS-2.1 models.
DeanoC
force-pushed
the
feat/hip-topk-argsort
branch
from
July 11, 2026 13:10
b3266cf to
3a481c5
Compare
The GPU top-K + log-prob extraction (geometric_draft_topk_cuda.cu) was compiled only on CUDA, so DFLASH_GPU_DRAFT_TOPK was a no-op on ROCm and AMD paid a per-step vocab x n_tokens D2H + CPU heap extract every speculation step. The kernel body is plain arithmetic (no tensor cores, no CUDA-only intrinsics), so the same .cu compiles unchanged for HIP. Compile geometric_draft_topk_cuda.cu directly with LANGUAGE HIP through the hip_compat <cuda_runtime.h> shim (the shared-.cu pattern already used by deepseek4_hc_cuda.cu), instead of adding a separate .hip.cu wrapper translation unit. Rename the backend guard macro DFLASH27B_HAVE_DRAFT_TOPK_CUDA -> DFLASH27B_HAVE_DRAFT_TOPK so the consumers (qwen35_dflash_target, test_dflash) are backend-neutral. test_draft_topk_cuda is now built on the HIP backend as well (CUDA spellings mapped by the same shim) and validated against the CPU reference on gfx1151 (Radeon 8060S, wave32, ROCm 6.4.4). Also qualify the README GPU-flag defaults to the server harness.
* test(deepseek4): add HC pre-mix CPU-parity test * test(deepseek4): scope test comments to hc_pre, flag n_hc>4 kernel limit * docs(deepseek4): correct hc_pre_mix scope note in HC parity test deepseek4_cuda_hc_pre_mix is called on the split mix plus CPU-finish decode path, so the 'currently uncalled anywhere in server/' note was wrong. Describe it as a separate mix-only entry instead. --------- Co-authored-by: mrciffa <49000955+davide221@users.noreply.github.com>
…rg#529) Summarize all DFLASH_*/DFLASH27B_* env vars into variables.md, grouped by subsystem, with a status legend marking debug-only, kill-switch, test/bench, and removal-candidate variables.
…l_mat_vec_q (Luce-Org#521) Wrap the cross-wave reduction epilogue of mul_mat_vec_q in if constexpr (nwarps > 1). For single-warp instantiations this removes a block-wide __syncthreads() and the shared-memory staging that the reduction would otherwise emit, without changing any computed result. The reduction is a no-op when there is only one warp: the shared array is only written under threadIdx.y > 0 (which never holds), the combine loop runs nwarps-1 == 0 iterations, and the barrier synchronizes a single warp. Guarding the block on nwarps > 1 makes that explicit so the compiler no longer has to emit the barrier for the single-wave decode path. For nwarps > 1 the statements are unchanged, only nested inside the guard.
…rence (Luce-Org#515) * perf(qwen35): add opt-in F32 rollback checkpoints * fix(qwen35): address rollback policy review findings * refactor(qwen35): centralize rollback diagnostics * fix(qwen35): guard null stream in rollback diagnostics Treat a null diagnostics stream as a no-op before calling fprintf, and cover the diagnostics-enabled null-stream path in the rollback policy test. Addresses the latest Cubic P3 on PR Luce-Org#515. --------- Co-authored-by: Drew Schuyler <dsky73@gmail.com>
* feat(ggml): vendor ROCmFPX quant support for DeepSeek V4 Flash Ports the ROCmFP4/ROCmFPX ggml work (previously lucebox-ggml Luce-Org#36, on the old submodule pointer) into the vendored server/deps/llama.cpp tree. - ROCmFP4/ROCmFPX quant types + CPU reference conversions (ggml/rocmfp4, ggml/rocmfpx) and the ggml type-trait registrations. - CUDA/HIP dequant, copy, getrows, MMVQ vecdot, MMVF, unary and FA paths for the new types. - Fused DS4 hyper-connection op (GGML_OP_DS4_HC) with the register-resident sinkhorn kernel; inert unless emitted by the DS4 fused-decode path. - DS4 SwiGLU split op plumbing for the fused FFN matvec paths. Layered on current main, so the main-side ggml work is preserved (fp64 RoPE reduction, Luce-Org#497 RDNA MMQ tile, LUCE_MMQ_DP_MAX_NE1, MMVQ_MAX_MOE_BATCH_SIZE, fused dual set_rows, raw-span guard). Review fixes on top of Luce-Org#36: - ggml_ftype_to_ggml_type: handle the 11 new ROCmFPX ftypes (dominant-type mapping) so the enum switch is -Wswitch/-Werror clean. - FP6 MMVQ vecdot: pad qs[] to avoid a stack over-read of the last window (bit-identical; the over-read bits were already masked out). * fix(ggml): make ROCmFP4 codebook self-contained on CUDA The non-HIP fallback of rocmfp4_get_int_from_codebook_16 / rocmfp4_get_low_int_from_codebook_16 called get_int_from_table_16, which is defined in vecdotq.cuh. TUs that pull in this header without vecdotq.cuh (fattn-chunked.cu reaches it via the fattn dequant chain) failed to compile under nvcc: rocmfp4_hip_codebook.cuh: error: identifier "get_int_from_table_16" is undefined The HIP path never hit this (it uses __builtin_amdgcn_perm), so the ROCm CI and the Strix build stayed green while the sm_86 CUDA build broke. Fix: inline the generic table expander (the generic branch of get_int_from_table_16, verbatim) as a static helper in this header, so the fallback no longer depends on include order. Bit-identical; the HIP hot path is unchanged. * fix(ggml): harden ROCmFP formats and DS4 ops
* fix(hip): use device wave width in RMSNorm Use HIP's target-specific warpSize builtin for the two-stage RMSNorm reduction. ROCm 7.2 no longer defines the deprecated AMDGCN wavefront-size preprocessor spellings used by the prior implementation, which prevented clean gfx1151 builds. The device builtin supports wave32 and wave64 without an architecture guess. * feat(deepseek4): add monolithic HIP decode path Add an explicit single-graph decode mode for single-device HIP deployments and a validated routed-expert top-k policy exposed through server CLI options. Keep the default model policy unchanged unless requested. The fused graph retains HC, attention, MoE, cache updates, and output projection on-device while caching structural compressor variants. Reset all request cache topology, handle raw-ring wraparound, apply expert limits to hash-routed layers, and expose full-graph timing. Add cache and ring regression coverage and document the serving contract. * perf(ggml): specialize ROCmFPX decode on gfx1151 Specialize the profiled single-column ROCmFP2/FP3 shapes with fixed-K loops, partially unroll the ROCmFP4-fast loop, and decode packed FP3 values with v_perm_b32 on HIP. Dispatch is restricted to gfx1151 and the exact validated shapes; all other devices and dimensions retain the generic kernels. The specializations preserve each lane's K traversal and accumulation order. * fix(deepseek4): prevent lost HC worker wakeups Publish the worker-pool generation and shutdown predicates under the same mutex used by the condition-variable wait. Without that synchronization, a worker could miss notify_all between checking the predicate and blocking, leaving the request thread spinning forever. Always notify after publishing and remove the racy sleeper-count optimization. * fix(deepseek4): honor fused decode on HIP
…3.5/RDNA4 The TOP_K path full-sorts every row via bitonic argsort regardless of k, so its latency is independent of k. Make it partial: - argsort.cu: cache keys in shared memory and bound the bitonic sort by the padded k, so work scales with k rather than ncols. - argsort.cu: GGML_ARGSORT_SHFL_XOR — for bitonic inner stages whose stride is below the wavefront width, both compare-exchange partners live in the same wave, so the swap is done register-to-register via __shfl_xor with no shared memory round-trip or __syncthreads() (wave32 on RDNA3.5/RDNA4). - top-k.cu: dedicated k=1 argmax kernel that skips the sort entirely. On HIP the CUB fast path is unavailable (GGML_CUDA_USE_CUB undefined), so the bitonic kernel is the hot path on AMD. test-backend-ops -o TOP_K passes on gfx1201 and gfx1151. Microbench A/B (ne=[1000,16]): k=1 15.59->2.59us (6.0x), k=10 ->4.58us (3.4x), k=400 ->6.56us (2.4x); geomean ~3.5-3.8x.
…ipCUB
The GPU argsort/top-k path uses a single-block bitonic sort (block_dims =
ncols_pad), so it caps at 1024 threads/block. On CUDA, ncols > 1024 is handled
by the CUB device-sort path; on HIP that path was compiled out (GGML_CUDA_USE_CUB
is CUDA-only), so ggml_backend_cuda_supports_op reported TOP_K/ARGSORT as
unsupported for ncols > 1024 and they fell back to the CPU reference — 275 of the
test-backend-ops cases (full-vocab sampling, large sorts) on every ROCm build.
Enable the existing device-sort path on HIP via hipCUB (rocPRIM), which provides
the same DeviceRadixSort / DeviceSegmentedSort / DeviceSegmentedRadixSort API:
* common.cuh: define GGML_CUDA_USE_HIPCUB on HIP.
* argsort.cu/.cuh: include <hipcub/hipcub.hpp> and compile argsort_f32_i32_cuda_cub
on HIP too; hipCUB has no CCCL strided iterator, so the init_offsets segment
path is used (already the CUB fallback).
* top-k.cu: on HIP, route ncols > 1024 through the device sort + slice-first-k,
mirroring the CUB branch; the fast partial-bitonic top-k still handles the
common ncols <= 1024 decode/verify case unchanged.
* ggml-cuda.cu: supports_op returns true for TOP_K/ARGSORT when either CUB or
hipCUB is available.
test-backend-ops -o TOP_K and -o ARGSORT: 0 not-supported / 0 fail on gfx1201
(R9700) and gfx1151 (Strix Halo) — the 275 previously-unsupported ncols>1024
cases (up to ne=[262144]) now run on-GPU and are bit-for-bit vs the CPU
reference. Requires the rocm hipcub + rocprim headers (ship with ROCm).
…nic + ds_swizzle (~1.19x)
Two schedule-level improvements to the ncols<=1024 partial-bitonic top-k, both
correctness-preserving (comparator/merge networks byte-identical → bit-exact
TOP_K set) and validated with interleaved A/B rebuilds on gfx1201 and gfx1151:
* Template the smallk / 2wave / shared bitonic kernels on kpad (and warp width)
with #pragma unroll, dispatched via switch(kpad) over the power-of-two ladder,
so the fixed-length shuffle/merge chains unroll fully and drop loop overhead.
Fully general — nothing keyed on the benchmark's ncols or k.
* Intra-wave xor shuffles (j < 32) issue a single ds_swizzle_b32 (bitmask mode,
(j<<10)|0x1f) instead of __shfl_xor, which on RDNA lowers to ds_bpermute_b32
and builds a per-lane address VGPR first. These kernels are latency-bound on
the dependent shuffle chains, so removing the address-compute shortens the
critical path. j==32 (wave64 tail) falls back to __shfl_xor.
test-backend-ops -o TOP_K: 0 fail on gfx1201 (R9700) and gfx1151 (Strix Halo).
Perf (ne=[1000,16], geomean over k=1/10/40/400 vs the prior partial-bitonic):
gfx1201 ~1.18x, gfx1151 ~1.19x; the k>1 cases gain 1.21-1.26x (k=1 argmax was
already at its floor and is unchanged).
…fles (~1.06-1.10x)
Follow-up tuning of the partial-bitonic top-k, correctness-preserving (comparator/
merge networks byte-identical -> bit-exact TOP_K set), validated with interleaved
A/B rebuilds on gfx1201 and gfx1151.
The intra-wave xor shuffles were issuing through the LDS permute crossbar
(ds_swizzle_b32 / ds_bpermute_b32). Move the ones expressible as pure VALU
cross-lane gathers off the crossbar onto the ALU pipe, shortening the dependent
shuffle chains these latency-bound kernels sit on:
* xor ^1 / ^2 (intra-quad) -> DPP mov_dpp quad_perm (0xB1 / 0x4E).
* xor ^16 (widest intra-wave stride, entry stage of every 32-lane chain) ->
v_permlanex16_b32 in its identity-cross form.
* ^4 / ^8 stay on ds_swizzle_b32; the wave64 j==32 tail falls back to __shfl_xor.
* plus a middle-cross-wave-stride group-merge fuse in the shared/2wave Phase-A/B.
All cross-lane encodings verified bit-identical to __shfl_xor on gfx1201 and
gfx1151. test-backend-ops -o TOP_K: 0 fail on both archs. Perf (ne=[1000,16],
geomean over k=1/10/40/400 vs the prior partial-bitonic): gfx1201 ~1.06x,
gfx1151 ~1.10x (k=1 and k=400 gain most: gfx1151 1.13x / 1.15x).
DeanoC
force-pushed
the
feat/hip-topk-argsort
branch
from
July 17, 2026 17:22
3a481c5 to
babb195
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
9 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Base is
base/hip-draft-topk-488(upstream Luce-Org#488's commit, @cheese-cakee's "compile DFlash GPU draft top-K for HIP"), so the diff is only our four commits.Optimizes the ggml
TOP_K/ARGSORTCUDA-HIP kernels for RDNA (CUB unavailable on HIP → the bitonic path is the hot path on AMD; this is where DSpark's indexer/draft/verify top-K lands).Impact — kernel microbench (
test-backend-ops perf -o TOP_K, ne=[1000,16])End-to-end — DeepSeek-V4-Flash + DSpark decode (gfx1151, spec-decode on)
Measured on the real workload (DS4 main + DSpark drafter,
dflash_server, decode of 200 tokens):+12% end-to-end decode. The win comes from the verify step (95.8 → 83.8 ms/step): the DSpark verify runs the indexer
top_k=512across all 43 layers (batched over the draft chain), so faster TOP_K/ARGSORT compounds. accept_rate unchanged (0.31) — top-K speed doesn't affect correctness/acceptance.Changes (4 commits, all correctness-preserving; comparator/merge networks byte-identical → bit-exact TOP_K set)
argsort.cu/top-k.cu: shared-mem key cache + partial bitonic bounded by padded k (work scales with k, not ncols);GGML_ARGSORT_SHFL_XORregister-to-register sub-wavefront compare-exchange; dedicated k=1 argmax;smallk/2wavespecializations.common.cuh/argsort.*/ggml-cuda.cu: enable the device-sort path on HIP via hipCUB (rocPRIM). Previouslysupports_opgated TOP_K/ARGSORT to ne[0]≤1024 on every ROCm build (275 test-backend-ops cases fell back to CPU); now they run on-GPU (cases up to ne=[262144] verified bit-exact vs CPU). Partial-bitonic still handles ncols≤1024 (~5× faster than hipCUB there).#pragma unroll+switch(kpad));ds_swizzle_b32for intra-wave xor shuffles.quad_perm(xor ^1/^2) +permlanex16(^16) instead of the LDS-crossbards_swizzle/ds_bpermute, shortening the dependent shuffle chains (small extra gain, ~+0.5% e2e).Validation
test-backend-ops test -o TOP_K/-o ARGSORT:2/2 backends passed, 0 fail, on gfx1201 (R9700) and gfx1151 (Strix Halo).Requires the ROCm
hipcub+rocprimheaders (ship with ROCm).