hugo/DS4Server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
6560f20129cfe729ea552d6000fb431cd62ebe2d
DS4Server/native/metal/ds4_gpu.h

2604 lines
102 KiB
C
Raw Blame History

#ifndef DS4_GPU_H
#define DS4_GPU_H
#include <stdbool.h>
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
/* =========================================================================
* GPU Tensor and Command Lifetime.
* =========================================================================
*
* Opaque device tensor used by the DS4-specific GPU executor.
*
* The public GPU API is tensor-resident: activations, KV state, and scratch
* buffers stay device-owned across the whole prefill/decode command sequence.
*/
#ifndef DS4_GPU_TENSOR_DEFINED
#define DS4_GPU_TENSOR_DEFINED
typedef struct ds4_gpu_tensor ds4_gpu_tensor;
#endif
#ifndef DS4_GPU_ATTENTION_DECODE_ROW_DEFINED
#define DS4_GPU_ATTENTION_DECODE_ROW_DEFINED
#define DS4_GPU_ATTENTION_DECODE_BATCH_MAX 32u
typedef struct {
uint64_t raw_kv;
uint64_t comp_kv;
uint64_t topk;
uint32_t pos;
uint32_t n_raw;
uint32_t raw_cap;
uint32_t raw_start;
uint32_t n_comp;
uint32_t top_k;
uint32_t window;
uint32_t ratio;
uint32_t indexed;
} ds4_gpu_attention_decode_row;
#endif
int ds4_gpu_init(void);
void ds4_gpu_cleanup(void);
ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes);
ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes);
ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes);
void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor);
uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor);
void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor);
int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count);
int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes);
int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes);
int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset,
const ds4_gpu_tensor *src, uint64_t src_offset,
uint64_t bytes);
int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset,
const ds4_gpu_tensor *src, uint64_t src_offset,
uint64_t count);
int ds4_gpu_moe_handoff_pack_tensor(
ds4_gpu_tensor *packed,
const ds4_gpu_tensor *ffn_norm,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_embd,
uint32_t n_expert);
int ds4_gpu_pack_slot_rows_f32_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *slots,
uint32_t n_rows,
uint32_t width,
uint32_t n_slots,
uint32_t slot_cap);
int ds4_gpu_begin_commands(void);
int ds4_gpu_flush_encoder(void);
int ds4_gpu_flush_commands(void);
int ds4_gpu_commands_active(void);
int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value);
int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label);
int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label);
#ifdef DS4_ROCM_BUILD
int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor,
uint64_t offset,
void *data,
uint64_t bytes,
uint64_t event_value,
const char *label);
#endif
int ds4_gpu_end_commands(void);
int ds4_gpu_synchronize(void);
int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size);
int ds4_gpu_set_model_fd(int fd);
int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map);
int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes);
int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes);
int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label);
int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label);
int ds4_gpu_q8_cache_suppressed(void);
void ds4_gpu_set_q8_cache_suppressed(int suppressed);
#ifdef DS4_ROCM_BUILD
void ds4_gpu_release_q8_f16_cache(void);
#endif
/* Model-file ranges assigned to CUDA devices by the multi-GPU placement
* planner. Metal keeps these declarations for the shared engine interface. */
#ifndef DS4_MAX_GPUS
#define DS4_MAX_GPUS 16
#endif
typedef struct {
uint64_t source_offset;
uint64_t bytes;
int target_device;
} ds4_tensor_range;
int ds4_gpu_device_cache_tensors(int device_id,
const ds4_tensor_range *ranges,
int n_ranges);
int ds4_gpu_register_support_map(const void *map, uint64_t size, uint64_t bias);
int ds4_gpu_device_cache_support_tensors(int device_id,
int entry_device_id,
const ds4_tensor_range *ranges,
int n_ranges,
int from_main_map);
uint64_t ds4_gpu_tier_free_vram(int logical_tier);
int ds4_gpu_lookup_cache(uint64_t source_offset, uint64_t bytes,
int *out_device_id, void **out_device_ptr);
int ds4_gpu_lookup_cache_device(uint64_t source_offset, uint64_t bytes);
int ds4_gpu_pro_q4_expert_table_auto_available(void);
int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size,
uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset,
uint64_t gate_expert_bytes, uint64_t down_expert_bytes,
uint32_t n_total_expert);
int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes);
void ds4_gpu_set_quality(bool quality);
void ds4_gpu_set_glm_model(bool enabled);
void ds4_gpu_set_ssd_streaming(bool enabled);
void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled);
#ifdef __APPLE__
void ds4_gpu_release_zero_prefix_prefill_mask_cache(void);
#endif
void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts);
void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes);
uint64_t ds4_gpu_recommended_working_set_size(void);
uint32_t ds4_gpu_stream_expert_cache_configured_count(void);
uint32_t ds4_gpu_stream_expert_cache_current_count(void);
typedef struct ds4_gpu_stream_expert_table {
const void *model_map;
uint64_t model_size;
uint32_t layer;
uint32_t n_total_expert;
uint64_t gate_offset;
uint64_t up_offset;
uint64_t down_offset;
uint64_t gate_expert_bytes;
uint64_t down_expert_bytes;
} ds4_gpu_stream_expert_table;
/* Reset only the prompt-local eviction heuristic. The resident SSD expert
* cache itself is intentionally kept warm across sessions. */
void ds4_gpu_stream_expert_cache_reset_route_hotness(void);
void ds4_gpu_stream_expert_cache_release_resident(void);
uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes);
int ds4_gpu_stream_expert_cache_seed_selected(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected);
int ds4_gpu_stream_expert_cache_begin_selected_load(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected);
int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
const ds4_gpu_stream_expert_table *table,
const ds4_gpu_tensor *selected,
uint32_t n_selected);
#ifdef __APPLE__
/* The async selected-load worker registers itself so Metal cache paths never
* wait on command buffers from that thread (they fail the load instead and
* the caller retries synchronously). */
void ds4_gpu_stream_expert_cache_note_service_thread(void);
#endif
#if defined(DS4_ROCM_BUILD) || (!defined(DS4_NO_GPU) && !defined(__APPLE__))
int ds4_gpu_stream_expert_cache_prepare_selected_batch(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_tokens,
uint32_t n_selected);
#endif
#ifdef DS4_ROCM_BUILD
int ds4_gpu_stream_expert_cache_load_layer(
const ds4_gpu_stream_expert_table *table);
int ds4_gpu_stream_expert_cache_seed_from_layer_selected(
const ds4_gpu_stream_expert_table *table,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_seed_tokens,
uint32_t n_selected);
int ds4_gpu_stream_expert_cache_release_layer_cache(void);
#endif
int ds4_gpu_stream_expert_cache_seed_experts(
const ds4_gpu_stream_expert_table *table,
const int32_t *expert_ids,
const uint32_t *expert_priorities,
uint32_t n_experts);
void ds4_gpu_print_memory_report(const char *label);
/* Tensor-parallel per-layer gates (Metal only). The encoder calls
* ds4_gpu_tp_gate_encode() right after the kernels that produce a partial
* block output in the TP slab: it closes the current encoder, makes the GPU
* signal a shared event, queues the exchange on a service thread, and makes
* the GPU wait for the CPU-signaled release before the combine kernel runs.
* Sequence values are assigned internally and increase monotonically; both
* ranks encode the identical gate sequence so values pair up by
* construction. The exchange callback runs on the service thread and must
* return nonzero on success. */
typedef int (*ds4_gpu_tp_exchange_fn)(void *ud, uint32_t layer, uint32_t gate, uint64_t seq);
/* Bind one rank of the two-way split. slab is the transport slab tensor and
* gpu_flags_off is the offset of its GPU-written gate-ready flag words. */
int ds4_gpu_tp_init(uint32_t rank,
ds4_gpu_tensor *slab, uint64_t gpu_flags_off,
ds4_gpu_tp_exchange_fn fn, void *ud);
void ds4_gpu_tp_shutdown(void);
/* Multi-session TP reuses slab slots across several encoded graph tapes.
* Shared-event arrival is required in that mode to make each partial vector
* CPU-visible before the transport thread reads it. */
void ds4_gpu_tp_set_session_batch_mode(int enabled);
/* The coordinator-only DSpark support model does not participate in TP.
* Suspend ownership only while encoding it; base-model verification remains
* split across both ranks. */
void ds4_gpu_tp_suspend_expert_sharding(int suspend);
int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate);
/* Verify-block batch gates: one exchange per layer moving `rows` partial
* rows at once (speculative verify). The callback runs on the gate service
* thread with the same ud as the row-gate exchange fn. */
typedef int (*ds4_gpu_tp_batch_exchange_fn)(void *ud, uint32_t layer,
uint32_t rows, uint64_t seq);
void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn);
int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows);
/* Prefill batch gates: the service thread exchanges `bytes` between two
* CPU-visible bounce tensors directly (payloads far beyond slab slots). */
typedef int (*ds4_gpu_tp_big_exchange_fn)(void *ud, uint32_t layer,
uint64_t seq, const void *out,
void *in, uint64_t bytes);
void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn);
int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows,
const ds4_gpu_tensor *out_t,
ds4_gpu_tensor *in_t,
uint64_t bytes);
/* Split big gate: kick publishes the GPU arrival marker (batch shared
* event, whose completion semantics make the bounce payload visible to
* the exchange thread) and queues the exchange, returning the gate seq
* (0 on failure); wait encodes the release. Multiple kicks may be in
* flight; waiting on the last seq covers all earlier kicks (monotonic
* release event, in-order service thread). */
uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows,
const ds4_gpu_tensor *out_t,
ds4_gpu_tensor *in_t,
uint64_t bytes);
int ds4_gpu_tp_big_gate_wait(uint64_t seq);
/* Pause/resume the DVFS keep-alive around work that keeps the GPU busy.
* No-op when TP is not bound. */
void ds4_gpu_tp_keepalive_pause(int paused);
/* Split attention heads across the two TP ranks in the GLM batch-prefill
* attention kernels (qk-low, attention-lora, value-project). The caller
* zeroes the unowned head range of the heads buffer and combines the
* attn-output partials over the TP big-gate exchange. */
void ds4_gpu_tp_set_attn_head_split(int enabled);
/* Skip the whole-file model residency set (TP sharding: only the
* owned ranges are warmed; the rest must never be paged in). Call before
* the model is mapped. */
void ds4_gpu_model_residency_skip(int skip);
/* Nonzero after any gate exchange failed; the eval must abort. */
int ds4_gpu_tp_failed(void);
/* Tensor-parallel sliced projections (Metal decode path only).
*
* ds4_gpu_matmul_q8_0_kslice_tensor computes a k-range partial matvec:
* out[out_dim] = W[:, k_off : k_off + k_cnt] @ x[x_elem_off : +k_cnt] where
* W rows span full_in_dim quantized Q8_0 elements. k offsets/counts must be
* multiples of 32 (Q8_0 block). Partial results from both ranks sum to the
* full projection.
*
* ds4_gpu_attention_output_q8_tp_tensor is the group-sliced attention output
* pair: low projection for groups [group0, group0+group_cnt) plus the
* matching k-slice of the expand projection, producing this rank's partial
* attention block output (n_tokens == 1 only). */
int ds4_gpu_matmul_q8_0_kslice_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t full_in_dim,
uint64_t k_off,
uint64_t k_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t x_elem_off);
/* CUDA multi-row variant. Each input row contains only the owned contiguous
* K slice, while each output row spans the full projection width. */
int ds4_gpu_matmul_q8_0_kslice_rows_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t full_in_dim,
uint64_t out_dim,
uint64_t k_off,
uint64_t k_cnt,
const ds4_gpu_tensor *x,
uint64_t n_rows);
int ds4_gpu_matmul_quant_kslice_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t full_in_dim,
uint64_t k_off,
uint64_t k_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t x_elem_off);
int ds4_gpu_attention_output_q8_tp_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *heads);
/* =========================================================================
* Embeddings and Indexer Helpers.
* =========================================================================
*
* These kernels seed HC state from token embeddings and implement the ratio-4
* compressed-attention indexer that chooses visible compressed rows.
*/
int ds4_gpu_embed_token_hc_tensor(
ds4_gpu_tensor *out_hc,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_embed_tokens_hc_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_embed_token_q8_0_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd);
int ds4_gpu_embed_tokens_q8_0_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd);
int ds4_gpu_embed_token_quant_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd);
int ds4_gpu_embed_tokens_quant_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd);
int ds4_gpu_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_head,
uint32_t head_dim,
float scale);
int ds4_gpu_indexer_scores_prefill_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale);
int ds4_gpu_indexer_scores_decode_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale);
int ds4_gpu_dspark_markov_argmax_tensor(ds4_gpu_tensor *out_idx,
const ds4_gpu_tensor *logits_row,
const void *model_map,
uint64_t model_size,
uint64_t w1_offset,
uint64_t w2_offset,
uint32_t prev_token,
uint32_t vocab,
uint32_t rank);
int ds4_gpu_indexer_topk_tensor(
ds4_gpu_tensor *selected,
const ds4_gpu_tensor *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k);
int ds4_gpu_indexer_top1_value_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *values,
const ds4_gpu_tensor *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t index_offset);
int ds4_gpu_matmul_q8_0_top1_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *values,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t index_offset);
int ds4_gpu_set_decode_fast_attention(int enabled);
int ds4_gpu_set_decode_score_vec4(int enabled);
/* GPU argmax over n_vocab F32 logits. Writes the winning index as int32 at
* out_idx[0]. Tie-break: lower index wins (matches host sample_argmax). */
int ds4_gpu_argmax_tensor(
ds4_gpu_tensor *out_idx,
const ds4_gpu_tensor *logits,
uint32_t n_vocab);
int ds4_gpu_dsv4_topk_mask_tensor(
ds4_gpu_tensor *mask,
const ds4_gpu_tensor *topk,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k);
/* =========================================================================
* Dense Projections, Norms, RoPE, and KV Rounding.
* =========================================================================
*
* The graph uses these primitives for Q/KV projections, HC/output projections,
* attention output projections, and DS4's tail-only RoPE.
*/
int ds4_gpu_matmul_q8_0_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_matmul_q8_0_decode_mpp_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_matmul_q8_0_rows_scalar_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_matmul_quant_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_matmul_quant_rows_scalar_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
/* Optional fused GPU operations.
*
* These are acceleration hooks, not required backend primitives. A backend
* that does not provide the fused kernel must still define the symbol and
* return 0. Callers then use the portable sequence of required primitives.
* Backends that return nonzero from a fused half-output operation must also
* implement the matching half-input HC expansion helpers below.
*/
int ds4_gpu_matmul_q8_0_pair_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
/* Multi-row decode projections that preserve the one-row reduction order. */
int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows);
int ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows);
int ds4_gpu_matmul_q8_0_f16_out_tensor(
ds4_gpu_tensor *out_h,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp);
int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor(
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *prequant,
uint32_t expert_split,
bool home_rank);
int ds4_gpu_shared_mid_swiglu_q8_0_tensor(
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp);
int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp);
int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float clamp);
int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float clamp);
int ds4_gpu_matmul_f16_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
/* Exact multi-row form of the DeepSeek 4096x256 F16 router projection. */
int ds4_gpu_matmul_f16_router_rows_exact_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
const ds4_gpu_tensor *x,
uint32_t n_rows);
int ds4_gpu_matmul_f16_pair_tensor(
ds4_gpu_tensor *out_a,
ds4_gpu_tensor *out_b,
const void *model_map,
uint64_t model_size,
uint64_t weight_a_offset,
uint64_t weight_b_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
/* Optional Metal decode fusion. Returns 1 when the paired projection and
* recurrent compressor-state store were encoded, 0 when the optimized path
* is unavailable, and -1 on an attempted-path error. */
int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos);
int ds4_gpu_matmul_f32_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok);
int ds4_gpu_repeat_hc_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *row,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_repeat_hc_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *rows,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_rms_norm_plain_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t n,
float eps);
int ds4_gpu_rms_norm_plain_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t n,
uint32_t rows,
float eps);
int ds4_gpu_rms_norm_weight_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
float eps);
int ds4_gpu_rms_norm_weight_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
uint32_t rows,
float eps);
int ds4_gpu_add_rms_norm_weight_tensor(
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *sum_out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
float eps);
int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_out,
const ds4_gpu_tensor *kv,
uint64_t kv_weight_offset,
uint32_t kv_n,
uint32_t rows,
float eps);
int ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_out,
const ds4_gpu_tensor *kv,
uint64_t kv_weight_offset,
uint32_t kv_n,
uint32_t rows,
uint32_t kv_n_head,
uint32_t kv_head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps);
int ds4_gpu_head_rms_norm_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
float eps);
int ds4_gpu_head_rms_norm_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps);
int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *q_half,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps);
int ds4_gpu_dsv4_fp8_kv_quantize_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t head_dim,
uint32_t n_rot);
int ds4_gpu_dsv4_indexer_qat_tensor(
ds4_gpu_tensor *x,
uint32_t n_rows,
uint32_t head_dim);
int ds4_gpu_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_kv_lora_rms_norm_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_raw,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
float eps);
int ds4_gpu_glm_k_b_project_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_norm,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head);
int ds4_gpu_glm_k_b_project_typed_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_norm,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head);
int ds4_gpu_glm_store_compact_kv_tensor(
ds4_gpu_tensor *kv_lora_cache,
ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *kv_norm,
const ds4_gpu_tensor *kv_raw,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
bool cache_f16);
int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_lora_cache,
ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *kv_raw,
uint64_t kv_weight_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
bool cache_f16,
float eps);
int ds4_gpu_glm_store_indexer_k_tensor(
ds4_gpu_tensor *indexer_key_cache,
const ds4_gpu_tensor *raw_k,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t bias_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t n_ctx_orig,
float eps,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16);
int ds4_gpu_glm_build_kv_cache_tensor(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16);
int ds4_gpu_glm_build_kv_cache_flash_tensor(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16);
int ds4_gpu_glm_attention_full_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16);
int ds4_gpu_glm_fill_selected_range_tensor(
ds4_gpu_tensor *selected,
uint32_t n_selected);
int ds4_gpu_glm_fill_selected_range_batch_tensor(
ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t pad_row);
int ds4_gpu_glm_indexer_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16);
int ds4_gpu_glm_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16);
int ds4_gpu_glm_qk_lowrank_q8_0_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim);
int ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim);
int ds4_gpu_glm_qk_lowrank_typed_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim);
int ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim);
int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *lora,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim);
int ds4_gpu_glm_value_project_typed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *lora,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim);
int ds4_gpu_glm_attention_indexed_decode_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_rope_tail_decode_rows_tensor(
ds4_gpu_tensor *x,
const ds4_gpu_attention_decode_row *rows,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_decode_typed_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor(
ds4_gpu_tensor *heads,
ds4_gpu_tensor *partial_lora,
ds4_gpu_tensor *partial_ms,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
bool selected_rows_valid,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
uint32_t block_rows,
uint32_t n_blocks,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(
ds4_gpu_tensor *heads,
ds4_gpu_tensor *partial_lora,
ds4_gpu_tensor *partial_ms,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
bool selected_rows_valid,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
uint32_t block_rows,
uint32_t n_blocks,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_batch_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_batch_typed_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_sort_i32_rows_asc_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint32_t row_width,
uint32_t n_rows);
int ds4_gpu_glm_attention_indexed_batch_lora_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
int ds4_gpu_glm_attention_flash_staged_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16);
int ds4_gpu_glm_attention_flash_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16);
/* Release decode fused KV finalizer: after the standalone RoPE kernel, this
* performs DS4's FP8 non-RoPE KV round trip and writes the F16-rounded raw
* attention cache row in one dispatch. */
int ds4_gpu_kv_fp8_store_raw_tensor(
ds4_gpu_tensor *kv,
ds4_gpu_tensor *raw_cache,
uint32_t raw_cap,
uint32_t row,
uint32_t head_dim,
uint32_t n_rot);
/* Exact multi-session form of the decode KV finalizer. KV rows are
* contiguous, while each output row is written to its session-private cache. */
int ds4_gpu_kv_fp8_store_raw_decode_rows_tensor(
ds4_gpu_tensor *kv,
ds4_gpu_tensor *const *raw_caches,
const uint32_t *raw_caps,
const uint32_t *raw_rows,
uint32_t n_rows,
uint32_t head_dim,
uint32_t n_rot);
/* Reference/raw-cache primitive kept for prefill and diagnostics. Decode uses
* ds4_gpu_kv_fp8_store_raw_tensor unless a diagnostic reference path is
* explicitly selected by the graph driver. */
int ds4_gpu_store_raw_kv_tensor(
ds4_gpu_tensor *raw_cache,
const ds4_gpu_tensor *kv,
uint32_t raw_cap,
uint32_t row,
uint32_t head_dim);
int ds4_gpu_store_raw_kv_batch_tensor(
ds4_gpu_tensor *raw_cache,
const ds4_gpu_tensor *kv,
uint32_t raw_cap,
uint32_t pos0,
uint32_t n_tokens,
uint32_t head_dim);
/* =========================================================================
* KV Compression and Attention.
* =========================================================================
*
* Compressed layers maintain rolling score/KV state and append pooled rows at
* ratio boundaries. Attention kernels consume raw SWA rows, compressed rows,
* and optional indexer masks.
*/
int ds4_gpu_compressor_update_tensor(
const ds4_gpu_tensor *kv_cur,
const ds4_gpu_tensor *sc_cur,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
ds4_gpu_tensor *comp_cache,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos,
uint32_t comp_row,
uint32_t n_rot,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps,
bool state_already_stored);
int ds4_gpu_compressor_store_batch_tensor(
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens);
int ds4_gpu_compressor_prefill_tensor(
ds4_gpu_tensor *comp_cache,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool quantize_fp8,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps);
int ds4_gpu_compressor_prefill_ratio4_replay_tensor(
ds4_gpu_tensor *comp_cache,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool quantize_fp8,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps);
int ds4_gpu_compressor_prefill_state_ratio4_tensor(
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv_tail,
const ds4_gpu_tensor *sc_tail,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t pos0);
int ds4_gpu_attention_decode_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_comp,
const ds4_gpu_tensor *comp_mask,
uint32_t use_mask,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_decode_heads_rope_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_comp,
const ds4_gpu_tensor *comp_mask,
uint32_t use_mask,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
int *fused_inv_rope);
/* Multi-session decode over contiguous Q/head rows and private KV caches.
* The row table is copied into CUDA launch parameters, so no device-side
* descriptor upload or synchronization is required. */
int ds4_gpu_attention_decode_rows_rope_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_attention_decode_row *rows,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
/* Diagnostic/public form of the dk=512 gathered decode-attention KV staging
* step. The compressed source must be F16; dst writes chronological raw-ring
* rows followed by compressed rows and must not overlap either source. */
int ds4_gpu_flash_kv_stage_f16_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_raw,
const ds4_gpu_tensor *comp,
uint32_t comp_is_f16,
uint32_t n_comp,
uint32_t head_dim);
int ds4_gpu_attention_prefill_raw_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t window,
uint32_t n_head,
uint32_t head_dim);
/* Rectangular raw prefill attention: q is a view of the n_q query rows at
* token positions [q_row0, q_row0 + n_q) of the chunk, raw_kv keeps all
* n_kv rows, heads receives n_q output rows. Used by the TP prefill row
* split; the square entry above is the q_row0 = 0, n_q = n_kv case. */
int ds4_gpu_attention_prefill_raw_heads_range_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t window,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_decode_raw_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t window,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_decode_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *topk,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t top_k,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_prefill_static_mixed_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim);
/* Rectangular static-mixed prefill attention: q is a view of the n_q query
* rows at token positions [q_row0, q_row0 + n_q) of the chunk, while raw_kv
* keeps all n_tokens rows and comp_kv all n_comp compressed keys. Used by
* the TP prefill row split; the square entry above is q_row0 = 0,
* n_q = n_tokens. */
int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_prefill_masked_mixed_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim);
int ds4_gpu_attention_output_q8_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp,
ds4_gpu_tensor *low_tmp,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens);
int ds4_gpu_attention_output_q4_K_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp,
ds4_gpu_tensor *low_tmp,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint32_t out_b_type,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens);
int ds4_gpu_attention_output_q8_batch_f16_tensor(
ds4_gpu_tensor *out_h,
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens);
int ds4_gpu_attention_output_low_q8_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
const ds4_gpu_tensor *heads);
int ds4_gpu_attention_output_low_q4_K_slice_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t group0,
uint32_t group_cnt,
const ds4_gpu_tensor *heads);
int ds4_gpu_attention_output_low_q8_rows_exact_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
const ds4_gpu_tensor *heads,
uint32_t n_rows);
int ds4_gpu_attention_output_q8_tp_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *heads);
/* =========================================================================
* Router, Shared Expert, and Routed MoE.
* =========================================================================
*
* These kernels implement the FFN body: router probabilities/top-k or hash
* routing, shared SwiGLU, and the IQ2_XXS/Q2_K/Q4_K routed experts.
*/
int ds4_gpu_swiglu_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *gate,
const ds4_gpu_tensor *up,
uint32_t n,
float clamp,
float weight);
int ds4_gpu_add_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
uint32_t n);
int ds4_gpu_add3_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const ds4_gpu_tensor *c,
uint32_t n);
int ds4_gpu_directional_steering_project_tensor(
ds4_gpu_tensor *x,
const ds4_gpu_tensor *directions,
uint32_t layer,
uint32_t width,
uint32_t rows,
float scale);
int ds4_gpu_router_select_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint64_t hash_offset,
uint32_t hash_rows,
uint32_t token,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_expert_groups,
uint32_t n_group_used,
bool has_bias,
bool hash_mode,
const ds4_gpu_tensor *logits);
int ds4_gpu_router_select_batch_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint64_t hash_offset,
uint32_t hash_rows,
uint32_t n_expert_groups,
uint32_t n_group_used,
bool has_bias,
bool hash_mode,
const ds4_gpu_tensor *logits,
const ds4_gpu_tensor *tokens,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens);
int ds4_gpu_glm_router_select_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
const ds4_gpu_tensor *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale);
int ds4_gpu_glm_router_select_batch_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
const ds4_gpu_tensor *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens);
int ds4_gpu_glm_routed_moe_one_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t layer_index,
const ds4_gpu_tensor *x,
bool force_resident);
int ds4_gpu_glm_routed_moe_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride,
bool force_resident);
int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride);
int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected);
void ds4_gpu_set_glm_mtp_verify_mode(bool enabled);
int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *block_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
uint64_t in_start,
uint64_t in_count,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_routed_moe_one_owned_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *experts,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t resident_expert_base,
uint32_t resident_expert_count,
float clamp,
const ds4_gpu_tensor *x,
ds4_gpu_tensor *down_output,
bool pack_fixed3,
ds4_gpu_tensor *shared_prequant);
int ds4_gpu_routed_moe_batch_owned_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *experts,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t resident_expert_base,
uint32_t resident_expert_count,
float clamp,
const ds4_gpu_tensor *x,
uint32_t layer_index,
uint32_t n_tokens,
bool *mid_is_f16);
int ds4_gpu_routed_moe_owned_slots_combine_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_slots,
const ds4_gpu_tensor *selected,
uint32_t out_dim,
uint32_t expert_split);
int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_slots,
const ds4_gpu_tensor *selected,
uint32_t out_dim,
uint32_t expert_split,
uint32_t rows);
int ds4_gpu_routed_moe_owned_packed_combine_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_packed,
const ds4_gpu_tensor *selected,
uint32_t out_dim,
uint32_t expert_split);
int ds4_gpu_routed_moe_one_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *experts,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float clamp,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *add_in,
uint32_t layer_index,
bool force_resident);
int ds4_gpu_routed_moe_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *experts,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float clamp,
const ds4_gpu_tensor *x,
uint32_t layer_index,
uint32_t n_tokens,
bool *mid_is_f16,
bool force_resident);
/* =========================================================================
* Hyper-Connection Kernels.
* =========================================================================
*
* HC kernels reduce four residual streams before a sublayer and expand the
* sublayer output back into four streams afterward.
*/
int ds4_gpu_hc_split_sinkhorn_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *mix,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps);
int ds4_gpu_hc_weighted_sum_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_weighted_sum_norm_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
const void *model_map,
uint64_t model_size,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
float norm_eps);
int ds4_gpu_hc_weighted_sum_split_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
/* Release decode fused HC pre-sublayer operation: split the HC mixer and
* immediately reduce four HC streams into the active 4096-wide sublayer row. */
int ds4_gpu_hc_split_weighted_sum_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps);
int ds4_gpu_hc_split_weighted_sum_norm_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps,
float norm_eps);
/* Batched HC RMSNorm followed by its narrow F16 mixer projection. On the
* tuned Metal path, scale_scratch stores one float per row instead of the
* full normalized HC tensor; other shapes retain the established fallback. */
int ds4_gpu_hc_rms_scale_project_f16_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *scale_scratch,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows,
float eps);
int ds4_gpu_output_hc_weights_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *pre,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_hc,
float eps);
int ds4_gpu_hc_expand_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *post,
const ds4_gpu_tensor *comb,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_expand_add_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *post,
const ds4_gpu_tensor *comb,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_expand_add_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *post,
const ds4_gpu_tensor *comb,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_expand_split_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_expand_split_half_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out_h,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_expand_add_split_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_expand_add_split_half_add_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add_h,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_shared_down_hc_expand_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *routed_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_shared_down_hc_expand_add_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *routed_out,
const ds4_gpu_tensor *routed_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_packed,
const ds4_gpu_tensor *selected,
uint32_t expert_split,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_matmul_q8_0_hc_expand_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *block_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc);
#ifdef __cplusplus
}
#endif
#endif
Reference in New Issue View Git Blame Copy Permalink