metal: optimize exact GLM Flash kernels and routing
Fuse HC, KDA, query projections and routed/shared stages; tile sparse attention and retain guarded selection fallbacks. Add state, boundary and bit-exact GPU oracles with model-appropriate vector fixtures.
This commit is contained in:
trueimage authored and 
jw committed
1 parent
5379c1c9ac
commit
f40230858b
19 files changed
+6606
-323
No files matched your search
@@ -9,6 +9,7 @@
|
||||
/ds4_test
|
||||
/ds4flash.gguf
|
||||
/TODO.md
|
||||
glm53flash_prefill*.md
|
||||
/gguf
|
||||
/core
|
||||
/core.*
|
||||
@@ -114,3 +115,8 @@ __pycache__/
|
||||
|
||||
/tests/test_qwen4_conv_parallel
|
||||
/tests/test_q8_prefill_variants
|
||||
|
||||
/tests/test_glm53_router_shared
|
||||
/tests/test_glm53_topk_fast
|
||||
/tests/test_glm53_q8_inputs
|
||||
/tests/test_glm53_hc_pre_repeat
|
||||
@@ -870,6 +870,28 @@ tests/test_image_decode.o: tests/test_image_decode.c ds4_image.h
|
||||
tests/test_image_decode: tests/test_image_decode.o ds4_image.o
|
||||
$(CC) $(CFLAGS) -o $@ $^ -lm
|
||||
|
||||
# Model-free exactness screens for the isolated M3 Ultra fork experiments.
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
GLM53_FORK_TESTS := tests/test_glm53_router_shared tests/test_glm53_topk_fast tests/test_glm53_q8_inputs
|
||||
$(GLM53_FORK_TESTS): %: %.c ds4_metal.o ds4_image.o ds4_gpu.h
|
||||
$(CC) $(CFLAGS) -I. $< ds4_metal.o ds4_image.o -o $@ $(METAL_LDLIBS)
|
||||
.PHONY: test-glm53-fork
|
||||
test-glm53-fork: $(GLM53_FORK_TESTS)
|
||||
./tests/test_glm53_router_shared
|
||||
./tests/test_glm53_topk_fast
|
||||
./tests/test_glm53_q8_inputs
|
||||
|
||||
tests/test_glm53_hc_pre_repeat: tests/test_glm53_hc_pre_repeat.c ds4.c ds4.h ds4_gpu.h ds4_metal.o ds4_image.o
|
||||
$(CC) $(CFLAGS) -Wno-unused-function -I. -ffunction-sections -fdata-sections $< ds4_metal.o ds4_image.o -Wl,-dead_strip -o $@ $(METAL_LDLIBS)
|
||||
|
||||
.PHONY: test-glm53-hc-pre-repeat
|
||||
test-glm53-hc-pre-repeat: tests/test_glm53_hc_pre_repeat
|
||||
./tests/test_glm53_hc_pre_repeat 0 0
|
||||
./tests/test_glm53_hc_pre_repeat 1 0
|
||||
./tests/test_glm53_hc_pre_repeat 0 1
|
||||
./tests/test_glm53_hc_pre_repeat 1 1
|
||||
endif
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_metal.o ds4_image.o
|
||||
$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
|
||||
@@ -878,6 +900,14 @@ $(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_cuda.o ds4_image.o $(MMQ_OBJS)
|
||||
$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
|
||||
endif
|
||||
|
||||
# Only the Metal build of this test is exercised by `make test`; the CUDA
|
||||
# variant still builds through `make test-glm53-kda`.
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
GLM53_KDA_DEFAULT_TEST := $(GLM53_KDA_TEST)
|
||||
else
|
||||
GLM53_KDA_DEFAULT_TEST :=
|
||||
endif
|
||||
|
||||
.PHONY: test-glm53-kda
|
||||
test-glm53-kda: $(GLM53_KDA_TEST)
|
||||
./$(GLM53_KDA_TEST)
|
||||
@@ -1283,7 +1313,8 @@ tests/test_web_recovery: tests/test_web_recovery.c ds4_web.c ds4_web.h
|
||||
|
||||
test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test test-session-state test-linux-memory test-engram test-web-recovery \
|
||||
tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \
|
||||
tests/test_deepseek4_vision_image tests/test_image_decode tests/test_prompt_prefix $(SAMPLING_TEST) ds4 ds4-server ds4-bench ds4-agent
|
||||
tests/test_deepseek4_vision_image tests/test_image_decode tests/test_prompt_prefix $(SAMPLING_TEST) $(GLM53_KDA_DEFAULT_TEST) \
|
||||
ds4 ds4-server ds4-bench ds4-agent
|
||||
./ds4-eval --validate-cases
|
||||
./ds4-eval --self-test-extractors
|
||||
./ds4_agent_test
|
||||
@@ -1296,6 +1327,7 @@ test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test test-session-
|
||||
./tests/test_sampling
|
||||
./tests/test_deepseek4_vision_image
|
||||
./tests/test_image_decode
|
||||
@if [ -n "$(GLM53_KDA_DEFAULT_TEST)" ]; then ./$(GLM53_KDA_TEST); fi
|
||||
|
||||
dspark-acceptance: ds4
|
||||
DS4_DSPARK_MODEL="$(DS4_DSPARK_MODEL)" \
|
||||
@@ -1351,6 +1383,7 @@ clean:
|
||||
rm -f tests/test_qwen4_ngrams
|
||||
rm -f tests/test_qwen4_ngram_state
|
||||
rm -f tests/test_web_recovery
|
||||
rm -f tests/test_glm53_hc_pre_repeat
|
||||
rm -f tests/test_metal_ssd_experts
|
||||
rm -f tests/test_metal_command_memory
|
||||
rm -f tests/test_deepseek41_metal
|
||||
|
||||
@@ -308,8 +308,28 @@ enum {
|
||||
DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT = 1u << 4,
|
||||
DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4 = 1u << 5,
|
||||
DS4_GPU_TEST_HC_RMS_SCALE_PROJ = 1u << 6,
|
||||
/* Exercise GLM tuning kernels on synthetic shapes/devices, while keeping
|
||||
* TP and streaming exclusions. The second flag forces the last KDA block
|
||||
* to finish before block 0 to test incoming-state ownership. */
|
||||
DS4_GPU_TEST_GLM53_PREFILL = 1u << 7,
|
||||
DS4_GPU_TEST_GLM53_KDA_LAST_BLOCK_FIRST = 1u << 8,
|
||||
};
|
||||
void ds4_gpu_test_set_flags(uint32_t flags);
|
||||
enum {
|
||||
DS4_GPU_GLM53_PREFILL_QK_LOW = 1u << 0,
|
||||
DS4_GPU_GLM53_PREFILL_INDEXED_ATTN = 1u << 1,
|
||||
DS4_GPU_GLM53_PREFILL_MOE_TAIL_CULL = 1u << 2,
|
||||
DS4_GPU_GLM53_PREFILL_KDA_PREPARE = 1u << 3,
|
||||
DS4_GPU_GLM53_PREFILL_KDA_RECURRENCE = 1u << 4,
|
||||
DS4_GPU_GLM53_ROUTER_TOP8 = 1u << 5,
|
||||
DS4_GPU_GLM53_DECODE_QK_LOW = 1u << 6,
|
||||
DS4_GPU_GLM53_BF16_SHORT_ROWS = 1u << 7,
|
||||
DS4_GPU_GLM53_DSA_SCORE_TILE = 1u << 8,
|
||||
DS4_GPU_GLM53_DECODE_KDA_VALUES4 = 1u << 9,
|
||||
DS4_GPU_GLM53_KDA_INPUTS = 1u << 10,
|
||||
};
|
||||
/* Returns and clears GLM tuning coverage recorded only in the test mode. */
|
||||
uint32_t ds4_gpu_test_glm53_prefill_take_dispatches(void);
|
||||
void ds4_gpu_release_zero_prefix_prefill_mask_cache(void);
|
||||
#else
|
||||
static inline int ds4_gpu_device_is_pre_m5_apple_silicon(void) { return 0; }
|
||||
@@ -1748,6 +1768,47 @@ int ds4_gpu_glm_attention_indexed_decode_typed_tensor(
|
||||
float beta_fast,
|
||||
float beta_slow);
|
||||
|
||||
int ds4_gpu_glm_attention_indexed_decode_exact_typed_tensor(
|
||||
ds4_gpu_tensor *heads,
|
||||
ds4_gpu_tensor *scores,
|
||||
ds4_gpu_tensor *lora,
|
||||
ds4_gpu_tensor *denom,
|
||||
const ds4_gpu_tensor *qk_low,
|
||||
const ds4_gpu_tensor *kv_lora_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);
|
||||
|
||||
int ds4_gpu_glm_attention_indexed_decode_exact_tensor(
|
||||
ds4_gpu_tensor *heads,
|
||||
ds4_gpu_tensor *scores,
|
||||
ds4_gpu_tensor *lora,
|
||||
ds4_gpu_tensor *denom,
|
||||
const ds4_gpu_tensor *qk_low,
|
||||
const ds4_gpu_tensor *kv_lora_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);
|
||||
|
||||
int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor(
|
||||
ds4_gpu_tensor *heads,
|
||||
ds4_gpu_tensor *partial_lora,
|
||||
@@ -3102,6 +3163,27 @@ int ds4_gpu_hc_expand_add_rms_norm_mix_split_norm_f16_tensor(
|
||||
float hc_eps,
|
||||
float norm_eps);
|
||||
|
||||
int ds4_gpu_hc_rms_norm_mix_split_norm_bf16_tensor(
|
||||
ds4_gpu_tensor *mix,
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *norm_out,
|
||||
ds4_gpu_tensor *split,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t mix_weight_offset,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint64_t norm_weight_offset,
|
||||
uint32_t n,
|
||||
uint32_t mix_dim,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t sinkhorn_iters,
|
||||
float eps,
|
||||
float hc_eps,
|
||||
float norm_eps);
|
||||
|
||||
#endif
|
||||
int ds4_gpu_output_hc_weights_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
@@ -3272,6 +3354,71 @@ int ds4_gpu_glm53_matmul_bf16_qkv(
|
||||
uint32_t out_dim,
|
||||
const ds4_gpu_tensor *x);
|
||||
|
||||
int ds4_gpu_glm53_matmul_bf16_pair(
|
||||
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,
|
||||
uint32_t in_dim,
|
||||
uint32_t out_dim,
|
||||
const ds4_gpu_tensor *x_a,
|
||||
const ds4_gpu_tensor *x_b);
|
||||
|
||||
#ifdef __APPLE__
|
||||
/* Single-token GLM-5.3 KDA input projections, all reading 4096 F32 values.
|
||||
* Slots are Q/K/V (8192 rows), f_a/g_a (128 rows), and beta (64 rows).
|
||||
* Returns zero without encoding when the resident M3 Ultra policy declines. */
|
||||
int ds4_gpu_glm53_kda_inputs_bf16(
|
||||
ds4_gpu_tensor *const outputs[6], const uint64_t weight_offsets[6],
|
||||
const void *model_map, uint64_t model_size, const ds4_gpu_tensor *x);
|
||||
|
||||
/* Exact bounded selector for serial GLM pool rows, with the original
|
||||
* sort as a GPU-dispatched fallback. Experimental opt-in. */
|
||||
int ds4_gpu_glm53_indexer_topk_tensor(ds4_gpu_tensor *selected,
|
||||
const ds4_gpu_tensor *scores, uint32_t n_comp,
|
||||
uint32_t n_tokens, uint32_t top_k);
|
||||
/* Call only after synchronizing. Reports cumulative selector telemetry. */
|
||||
int ds4_gpu_test_glm53_topk_stats(uint32_t *out, uint32_t count);
|
||||
int ds4_gpu_test_glm53_topk_poison_recovery_state(void);
|
||||
void ds4_gpu_test_invalidate_completion_counters(void);
|
||||
|
||||
int ds4_gpu_glm53_kda_inputs_q8_bf16(
|
||||
ds4_gpu_tensor *const outputs[6], const uint64_t weight_offsets[6],
|
||||
const void *model_map, uint64_t model_size, const ds4_gpu_tensor *x);
|
||||
#endif
|
||||
|
||||
uint64_t ds4_gpu_encoder_count(void);
|
||||
|
||||
int ds4_gpu_glm53_matmul_bf16_trio(
|
||||
ds4_gpu_tensor *out_a,
|
||||
ds4_gpu_tensor *out_b,
|
||||
ds4_gpu_tensor *out_c,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_a_offset,
|
||||
uint64_t weight_b_offset,
|
||||
uint64_t weight_c_offset,
|
||||
uint32_t in_dim,
|
||||
uint32_t out_dim_ab,
|
||||
uint32_t out_dim_c,
|
||||
const ds4_gpu_tensor *x);
|
||||
|
||||
int ds4_gpu_glm53_matmul_bf16_hc_expand4(
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *hc_out,
|
||||
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,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *post,
|
||||
const ds4_gpu_tensor *comb,
|
||||
uint32_t n_hc);
|
||||
|
||||
#ifndef DS4_GLM53_VISION_TYPES_DEFINED
|
||||
#define DS4_GLM53_VISION_TYPES_DEFINED
|
||||
#define DS4_GLM53_VISION_LAYERS 24u
|
||||
@@ -3668,6 +3815,17 @@ int ds4_gpu_qwen4_mtp_stage_tensor(
|
||||
int ds4_gpu_qwen4_mtp_combine_tensor(
|
||||
ds4_gpu_tensor *R_out, const ds4_gpu_tensor *proj, uint32_t n_embd, uint32_t n_hc);
|
||||
|
||||
#ifdef __APPLE__
|
||||
int ds4_gpu_glm53_router_shared_exact(
|
||||
ds4_gpu_tensor *logits, ds4_gpu_tensor *selected,
|
||||
ds4_gpu_tensor *weights, ds4_gpu_tensor *probs,
|
||||
ds4_gpu_tensor *counter, ds4_gpu_tensor *shared_mid,
|
||||
const void *model_map, uint64_t model_size,
|
||||
uint64_t router_offset, uint64_t bias_offset,
|
||||
uint64_t gate_offset, uint64_t up_offset,
|
||||
const ds4_gpu_tensor *x, float scale, float clamp);
|
||||
#endif
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
+1334
-58
File diff suppressed because it is too large.
Load diff
@@ -318,3 +318,335 @@ kernel void kernel_argsort_merge_f32_i32(
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC>;
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc_causal")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC, true>;
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc_causal_prefix")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC, true, true>;
|
||||
|
||||
// Adapted from IngeniousIdiocy/ds4 95eb218 (MIT).
|
||||
// Bounded exact DSA selector.
|
||||
// LEVER topk-select-fast: bounded radix fast path for the GLM-5.3 serial-decode
|
||||
// DSA pool selection.
|
||||
//
|
||||
// WHAT THE PRODUCTION CHAIN EMITS, AND WHY THIS CAN REPRODUCE IT
|
||||
// The block sort above is a comparison-exchange network whose predicate is
|
||||
// `lo.val < hi.val` with out-of-range indices forced last; the fused merge
|
||||
// takes the left run on `p0.val >= p1.val`. On an input whose valid scores
|
||||
// are all finite and whose 512 largest are pairwise distinct AND strictly
|
||||
// greater than every other valid score, that whole chain is a TOTAL ORDER BY
|
||||
// VALUE ALONE: each block emits its own top-512 in descending order, no global
|
||||
// winner is lost to the per-block truncation (a block cannot hold more than
|
||||
// 512 of the global top 512), and merging descending runs of distinct values
|
||||
// yields the descending union. Index order, block partition, chunking and the
|
||||
// NaN-dependent merge-path partition never enter. That is the accepted class,
|
||||
// and it is exactly this path's acceptance predicate; ANY other input takes the
|
||||
// unchanged production chain via an indirect dispatch whose threadgroup count
|
||||
// this path writes.
|
||||
//
|
||||
// KEY. bits ^ (bits>>31 ? 0xFFFFFFFF : 0x80000000) is order-isomorphic to the
|
||||
// float order over the full signed range. It would separate -0 from +0, which
|
||||
// the production comparisons treat as EQUAL, so -0 is normalised to +0 first;
|
||||
// a -0/+0 pair then shows up as an exact key tie and is rejected like any other
|
||||
// tie. Subnormals need no special handling: they are ordinary finite values
|
||||
// with a monotone bit pattern and the comparator compares them normally.
|
||||
// Non-finite values are classified by exponent bits over the COMPLETE valid
|
||||
// input in pass 1 (nothing is skipped by the narrowing) and force fallback.
|
||||
// SUBNORMALS. The shader library is built with MTLCompileOptions defaults,
|
||||
// i.e. fastMathEnabled = YES (ds4_metal.m ~7232; only DS4_METAL_MATH_SAFE, off
|
||||
// by default, pins MTLMathModeSafe). Under fast math the shipped comparator's
|
||||
// `<` / `>=` may flush denormals to zero, which this bit-level key does not.
|
||||
// Rather than guess which way the shipped comparator resolves a denormal, any
|
||||
// subnormal anywhere in the valid input is treated as an unsupported class and
|
||||
// takes the production chain. Zero subnormals occurred in 88 M captured real
|
||||
// scores. The classification is a pure integer exponent/mantissa test, so
|
||||
// fast math cannot alter it.
|
||||
// ===========================================================================
|
||||
|
||||
struct ds4_metal_args_glm53_topk_fast {
|
||||
uint32_t n_comp; // valid scores in the row
|
||||
uint32_t top_k; // winners to emit (512 for the DSA pool select)
|
||||
uint32_t hist_bits; // radix digit width for the single narrowing pass
|
||||
uint32_t cand_cap; // candidate buffer capacity (entries)
|
||||
uint32_t fb_count; // fallback dispatches whose grid this path gates
|
||||
uint32_t pad0;
|
||||
uint32_t fb_grid[16]; // (x,y) per fallback dispatch, in encode order
|
||||
};
|
||||
|
||||
#define DS4_TOPK_FAST_FLAG_NONFINITE 1u // Inf, NaN or subnormal
|
||||
#define DS4_TOPK_FAST_FLAG_OVERFLOW 2u
|
||||
#define DS4_TOPK_FAST_FLAG_NOBIN 4u
|
||||
#define DS4_TOPK_FAST_FLAG_TIE 8u
|
||||
#define DS4_TOPK_FAST_FLAG_SHORT 16u
|
||||
#define DS4_TOPK_FAST_FLAG_BADIDX 32u
|
||||
|
||||
// ctrl[0] fallback flags, ctrl[1] candidate counter, ctrl[2] accepted (1/0),
|
||||
// ctrl[3] candidate count observed by the finisher, ctrl[4] calls, ctrl[5]
|
||||
// accepted calls. ctrl[4]/ctrl[5] are cumulative telemetry, never cleared here.
|
||||
//
|
||||
// Cumulative rejection telemetry (added 2026-09-05;
|
||||
// none of it is cleared per call, and none of it feeds any decision):
|
||||
// ctrl[6] reject_flags_or - OR of the reject word over EVERY rejected call
|
||||
// in the process. It is NOT "the last rejection";
|
||||
// the old label `last_reject_flags` was wrong.
|
||||
// ctrl[7] current consecutive-rejection run length
|
||||
// ctrl[8] rejections whose reject word had NONFINITE (bit 0)
|
||||
// ctrl[9] ... OVERFLOW (bit 1)
|
||||
// ctrl[10] ... NOBIN (bit 2)
|
||||
// ctrl[11] ... TIE (bit 3)
|
||||
// ctrl[12] ... SHORT (bit 4)
|
||||
// ctrl[13] ... BADIDX (bit 5)
|
||||
// ctrl[14] longest consecutive-rejection run observed
|
||||
// A call can set more than one cause, so the six per-cause counters sum to at
|
||||
// least the rejected-call count; they are per-cause frequencies, not a
|
||||
// partition. The finisher is a single threadgroup and the selector calls are
|
||||
// serialised by dispatch order, so lane 0 owns this state.
|
||||
|
||||
static inline uint ds4_topk_fast_key(float v) {
|
||||
uint b = as_type<uint>(v);
|
||||
if (b == 0x80000000u) b = 0u; // -0 == +0 for the comparator
|
||||
return (b >> 31) ? ~b : (b ^ 0x80000000u);
|
||||
}
|
||||
|
||||
kernel void kernel_glm53_topk_fast_hist(
|
||||
constant ds4_metal_args_glm53_topk_fast & args,
|
||||
device const float * scores,
|
||||
device atomic_uint * ctrl,
|
||||
device atomic_uint * hist,
|
||||
threadgroup atomic_uint * lhist [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
uint3 tgpg [[threadgroups_per_grid]],
|
||||
ushort3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const uint nth = ntg.x;
|
||||
const uint tid = tpitg.x;
|
||||
const uint nbins = 1u << args.hist_bits;
|
||||
const uint shift = 32u - args.hist_bits;
|
||||
|
||||
for (uint b = tid; b < nbins; b += nth)
|
||||
atomic_store_explicit(&lhist[b], 0u, memory_order_relaxed);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
uint nonfinite = 0u;
|
||||
const uint stride = nth * tgpg.x;
|
||||
for (uint i = tgpig.x * nth + tid; i < args.n_comp; i += stride) {
|
||||
const float v = scores[i];
|
||||
const uint b = as_type<uint>(v);
|
||||
const uint e = b & 0x7f800000u;
|
||||
const uint m = b & 0x007fffffu;
|
||||
if (e == 0x7f800000u) nonfinite = 1u; // Inf or NaN
|
||||
if (e == 0u && m != 0u) nonfinite = 1u; // subnormal (see above)
|
||||
atomic_fetch_add_explicit(&lhist[ds4_topk_fast_key(v) >> shift],
|
||||
1u, memory_order_relaxed);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint b = tid; b < nbins; b += nth) {
|
||||
const uint c = atomic_load_explicit(&lhist[b], memory_order_relaxed);
|
||||
if (c) atomic_fetch_add_explicit(&hist[b], c, memory_order_relaxed);
|
||||
}
|
||||
if (nonfinite)
|
||||
atomic_fetch_or_explicit(&ctrl[0], DS4_TOPK_FAST_FLAG_NONFINITE,
|
||||
memory_order_relaxed);
|
||||
if (tgpig.x == 0u && tid == 0u) {
|
||||
atomic_fetch_add_explicit(&ctrl[4], 1u, memory_order_relaxed);
|
||||
if (args.n_comp < args.top_k)
|
||||
atomic_fetch_or_explicit(&ctrl[0], DS4_TOPK_FAST_FLAG_SHORT,
|
||||
memory_order_relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/* Every threadgroup redundantly reduces the same integer histogram to the same
|
||||
* boundary bin, which is bit-identical by construction (§3.3: redundant
|
||||
* recomputation in identical structure). That removes a dependent 1-TG scan
|
||||
* dispatch between the histogram and the gather. */
|
||||
kernel void kernel_glm53_topk_fast_gather(
|
||||
constant ds4_metal_args_glm53_topk_fast & args,
|
||||
device const float * scores,
|
||||
device atomic_uint * ctrl,
|
||||
device const uint * hist,
|
||||
device uint2 * cand,
|
||||
threadgroup uint * sscan [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
uint3 tgpg [[threadgroups_per_grid]],
|
||||
ushort3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const uint nth = ntg.x;
|
||||
const uint tid = tpitg.x;
|
||||
const uint nbins = 1u << args.hist_bits;
|
||||
const uint shift = 32u - args.hist_bits;
|
||||
const uint per = (nbins + nth - 1u) / nth; // bins owned per thread
|
||||
|
||||
// sscan[0..nth) : per-thread sum of its contiguous bin block
|
||||
// sscan[nth] : boundary bin, sscan[nth+1] : count strictly above it
|
||||
uint mine = 0u;
|
||||
const uint b0 = tid * per;
|
||||
for (uint b = b0; b < b0 + per && b < nbins; b++) mine += hist[b];
|
||||
sscan[tid] = mine;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// inclusive SUFFIX sum over sscan[0..nth) (Hillis-Steele, reversed)
|
||||
for (uint off = 1u; off < nth; off <<= 1) {
|
||||
const uint add = (tid + off < nth) ? sscan[tid + off] : 0u;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
sscan[tid] += add;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
if (tid == 0u) { sscan[nth] = 0xffffffffu; sscan[nth + 1u] = 0u; }
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// count strictly above this thread's block = suffix sum of the next block
|
||||
const uint above_block = (tid + 1u < nth) ? sscan[tid + 1u] : 0u;
|
||||
if (above_block < args.top_k && b0 < nbins) {
|
||||
uint acc = above_block;
|
||||
const uint hi = min(b0 + per, nbins);
|
||||
for (uint b = hi; b-- > b0; ) {
|
||||
const uint c = hist[b];
|
||||
if (acc < args.top_k && acc + c >= args.top_k) {
|
||||
sscan[nth] = b;
|
||||
sscan[nth + 1u] = acc;
|
||||
break;
|
||||
}
|
||||
acc += c;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint bin_lo = sscan[nth];
|
||||
if (bin_lo == 0xffffffffu) {
|
||||
if (tid == 0u)
|
||||
atomic_fetch_or_explicit(&ctrl[0], DS4_TOPK_FAST_FLAG_NOBIN,
|
||||
memory_order_relaxed);
|
||||
return;
|
||||
}
|
||||
|
||||
const uint stride = nth * tgpg.x;
|
||||
for (uint i = tgpig.x * nth + tid; i < args.n_comp; i += stride) {
|
||||
const uint key = ds4_topk_fast_key(scores[i]);
|
||||
if ((key >> shift) >= bin_lo) {
|
||||
const uint slot = atomic_fetch_add_explicit(&ctrl[1], 1u,
|
||||
memory_order_relaxed);
|
||||
if (slot < args.cand_cap) cand[slot] = uint2(key, i);
|
||||
else atomic_fetch_or_explicit(&ctrl[0], DS4_TOPK_FAST_FLAG_OVERFLOW,
|
||||
memory_order_relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* One threadgroup. Sorts the narrowed candidates descending by (key, index) -
|
||||
* a total order, so the result is deterministic regardless of the gather's
|
||||
* atomic arrival order - checks the acceptance predicate, and on acceptance
|
||||
* writes the ordered pool list. It also writes the fallback dispatches'
|
||||
* threadgroup counts (zero when accepted) and clears the scratch for the next
|
||||
* call. */
|
||||
kernel void kernel_glm53_topk_fast_finish(
|
||||
constant ds4_metal_args_glm53_topk_fast & args,
|
||||
device atomic_uint * ctrl,
|
||||
device uint * hist,
|
||||
device const uint2 * cand,
|
||||
device int32_t * out_idx,
|
||||
device uint32_t * indirect,
|
||||
threadgroup uint2 * shmem [[threadgroup(0)]],
|
||||
ushort3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const uint nth = ntg.x; // == args.cand_cap, one candidate per lane
|
||||
const uint tid = tpitg.x;
|
||||
|
||||
const uint flags = atomic_load_explicit(&ctrl[0], memory_order_relaxed);
|
||||
const uint count = atomic_load_explicit(&ctrl[1], memory_order_relaxed);
|
||||
const uint have = min(count, args.cand_cap);
|
||||
|
||||
/* Register-resident bitonic sort, descending by (key, index). Exactly the
|
||||
* structure of kernel_argsort_f32_i32_desc_pair above: every stage with
|
||||
* j < 32 is a simd_shuffle_xor and the rest use double-buffered threadgroup
|
||||
* staging, so 55 stages cost 15 barriers instead of 55. (key, index) is a
|
||||
* total order, so the result does not depend on the order in which the
|
||||
* gather's atomic counter handed out candidate slots. */
|
||||
uint2 cur = (tid < have) ? cand[tid] : uint2(0u, 0xffffffffu);
|
||||
threadgroup uint2 *buf0 = shmem;
|
||||
threadgroup uint2 *buf1 = shmem + nth;
|
||||
bool use_buf1 = false;
|
||||
for (uint k = 2u; k <= nth; k <<= 1) {
|
||||
for (uint j = k >> 1; j > 0u; j >>= 1) {
|
||||
const uint partner = tid ^ j;
|
||||
uint2 other;
|
||||
if (j >= 32u) {
|
||||
threadgroup uint2 *buf = use_buf1 ? buf1 : buf0;
|
||||
buf[tid] = cur;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
other = buf[partner];
|
||||
use_buf1 = !use_buf1;
|
||||
} else {
|
||||
other.x = simd_shuffle_xor(cur.x, (ushort)j);
|
||||
other.y = simd_shuffle_xor(cur.y, (ushort)j);
|
||||
}
|
||||
const bool cur_greater = (cur.x != other.x) ? (cur.x > other.x)
|
||||
: (cur.y < other.y);
|
||||
const bool want_greater = ((tid & k) == 0u) == (tid < partner);
|
||||
if (cur_greater != want_greater) cur = other;
|
||||
}
|
||||
}
|
||||
buf0[tid] = cur;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
threadgroup uint2 *s = buf0;
|
||||
|
||||
/* Acceptance predicate over the ordered candidates.
|
||||
*
|
||||
* REPAIRED 2026-09-05. The previous version had
|
||||
* every lane write an ordinary threadgroup `uint bad` (1 for a tie, 2 for
|
||||
* an invalid index). Those conflicting non-atomic stores are a data race
|
||||
* on a value that decides acceptance, the output-writing branch AND the
|
||||
* indirect fallback grids; a barrier around them does not serialise them
|
||||
* with each other, and the two causes overwrote one another so only one
|
||||
* could be reported. It is now a threadgroup atomic OR of DISTINCT cause
|
||||
* bits, reduced before the barrier that publishes it, so every lane reads
|
||||
* ONE uniform value and both causes survive. The bits are the same
|
||||
* DS4_TOPK_FAST_FLAG_* values used device-side, so the OR composes with
|
||||
* `flags` directly. No floating-point operation changes. */
|
||||
threadgroup atomic_uint bad_bits;
|
||||
if (tid == 0u) atomic_store_explicit(&bad_bits, 0u, memory_order_relaxed);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint ncheck = min(count, args.top_k + 1u); // covers the cut tie at 511/512
|
||||
uint mine = 0u;
|
||||
for (uint i = tid; i + 1u < ncheck && i + 1u < nth; i += nth)
|
||||
if (s[i].x == s[i + 1u].x) mine |= DS4_TOPK_FAST_FLAG_TIE;
|
||||
for (uint i = tid; i < args.top_k && i < nth; i += nth)
|
||||
if (s[i].y >= args.n_comp) mine |= DS4_TOPK_FAST_FLAG_BADIDX;
|
||||
if (mine) atomic_fetch_or_explicit(&bad_bits, mine, memory_order_relaxed);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
/* one uniform acceptance value, published after the reduction */
|
||||
uint reject = flags | atomic_load_explicit(&bad_bits, memory_order_relaxed);
|
||||
if (count < args.top_k || count > args.cand_cap || nth != args.cand_cap)
|
||||
reject |= DS4_TOPK_FAST_FLAG_SHORT;
|
||||
const bool accept = (reject == 0u);
|
||||
|
||||
if (accept) {
|
||||
for (uint i = tid; i < args.top_k; i += nth)
|
||||
out_idx[i] = (int32_t)s[i].y;
|
||||
}
|
||||
|
||||
// fallback dispatch grids: zero when accepted, the real grid otherwise
|
||||
for (uint d = tid; d < args.fb_count && d < 8u; d += nth) {
|
||||
indirect[d * 3u + 0u] = accept ? 0u : args.fb_grid[d * 2u + 0u];
|
||||
indirect[d * 3u + 1u] = accept ? 0u : args.fb_grid[d * 2u + 1u];
|
||||
indirect[d * 3u + 2u] = accept ? 0u : 1u;
|
||||
}
|
||||
|
||||
// clear the scratch for the next call
|
||||
const uint nbins = 1u << args.hist_bits;
|
||||
for (uint b = tid; b < nbins; b += nth) hist[b] = 0u;
|
||||
if (tid == 0u) {
|
||||
atomic_store_explicit(&ctrl[0], 0u, memory_order_relaxed);
|
||||
atomic_store_explicit(&ctrl[1], 0u, memory_order_relaxed);
|
||||
atomic_store_explicit(&ctrl[2], accept ? 1u : 0u, memory_order_relaxed);
|
||||
atomic_store_explicit(&ctrl[3], count, memory_order_relaxed);
|
||||
if (accept) {
|
||||
atomic_fetch_add_explicit(&ctrl[5], 1u, memory_order_relaxed);
|
||||
atomic_store_explicit(&ctrl[7], 0u, memory_order_relaxed); // run ends
|
||||
} else {
|
||||
atomic_fetch_or_explicit(&ctrl[6], reject, memory_order_relaxed);
|
||||
for (uint b = 0u; b < 6u; b++)
|
||||
if (reject & (1u << b))
|
||||
atomic_fetch_add_explicit(&ctrl[8u + b], 1u, memory_order_relaxed);
|
||||
const uint run = atomic_load_explicit(&ctrl[7], memory_order_relaxed) + 1u;
|
||||
atomic_store_explicit(&ctrl[7], run, memory_order_relaxed);
|
||||
if (run > atomic_load_explicit(&ctrl[14], memory_order_relaxed))
|
||||
atomic_store_explicit(&ctrl[14], run, memory_order_relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
+11
-6
@@ -70,7 +70,7 @@ struct ds4_metal_args_mul_mv_ext {
|
||||
int16_t r3;
|
||||
};
|
||||
|
||||
template<short NR0>
|
||||
template<short NR0, bool COHERENT_STORE = false>
|
||||
static inline void helper_mv_reduce_and_write(
|
||||
device float * dst_f32,
|
||||
float sumf[NR0],
|
||||
@@ -107,7 +107,12 @@ static inline void helper_mv_reduce_and_write(
|
||||
float tot = simd_sum(shmem_f32[row][tiisg]);
|
||||
|
||||
if (tiisg == 0 && sgitg == 0) {
|
||||
dst_f32[r0 + row] = tot;
|
||||
if (COHERENT_STORE) {
|
||||
atomic_store_explicit((device atomic_uint *)(dst_f32 + r0 + row),
|
||||
as_type<uint>(tot), memory_order_relaxed);
|
||||
} else {
|
||||
dst_f32[r0 + row] = tot;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -473,7 +478,7 @@ kernel void kernel_mul_mv_q8_0_f32_pair(
|
||||
// same lane that owns the reduced output row. The point is not to fuse two
|
||||
// independent weight streams into one matmul; it is to remove the separate
|
||||
// activation pass and its reread of the two 2048-wide rows.
|
||||
template<short NR0, bool STORE_GATE_UP>
|
||||
template<short NR0, bool STORE_GATE_UP, short NSG_OVERRIDE = 0>
|
||||
void kernel_dsv4_shared_gate_up_swiglu_q8_0_impl(
|
||||
constant ds4_metal_args_mul_mv & args,
|
||||
device const char * src0_gate,
|
||||
@@ -487,7 +492,7 @@ void kernel_dsv4_shared_gate_up_swiglu_q8_0_impl(
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
const short NSG = FC_mul_mv_nsg;
|
||||
const short NSG = NSG_OVERRIDE ? NSG_OVERRIDE : FC_mul_mv_nsg;
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
constexpr short NQ = 8;
|
||||
|
||||
@@ -1244,7 +1249,7 @@ typedef decltype(kernel_mul_mv_t_t<half, half>) mul_mv_t_t;
|
||||
template [[host_name("kernel_mul_mv_f32_f32")]] kernel mul_mv_t_t kernel_mul_mv_t_t<float, float>;
|
||||
template [[host_name("kernel_mul_mv_f16_f32")]] kernel mul_mv_t_t kernel_mul_mv_t_t<half, float>;
|
||||
|
||||
template<typename T0, typename T04, typename T1, typename T14, short NR0, typename args_t>
|
||||
template<typename T0, typename T04, typename T1, typename T14, short NR0, typename args_t, bool COHERENT_STORE = false>
|
||||
void kernel_mul_mv_t_t_4_impl(
|
||||
args_t args,
|
||||
device const char * src0,
|
||||
@@ -1322,7 +1327,7 @@ void kernel_mul_mv_t_t_4_impl(
|
||||
|
||||
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
||||
|
||||
helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
||||
helper_mv_reduce_and_write<NR0, COHERENT_STORE>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
||||
}
|
||||
|
||||
template<typename T0, typename T04, typename T1, typename T14, typename args_t>
|
||||
|
||||
+135
-54
@@ -1321,6 +1321,24 @@ kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2(
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/* Self-contained on purpose. glm53_bf16.metal has an identical helper, but
|
||||
* depending on it would couple this file to that one across the concatenated
|
||||
* library: pointing DS4_METAL_GLM53_BF16_SOURCE at an older revision of that
|
||||
* file would then stop THIS file compiling, which defeats the per-file source
|
||||
* overrides used for shader A/B runs. */
|
||||
static inline float4 ds4_hc_bf16x4_to_f32x4(ushort4 v) {
|
||||
return float4(as_type<float>((uint)v.x << 16),
|
||||
as_type<float>((uint)v.y << 16),
|
||||
as_type<float>((uint)v.z << 16),
|
||||
as_type<float>((uint)v.w << 16));
|
||||
}
|
||||
|
||||
static inline float4 ds4_hc_mix_widen(half4 v) { return float4(v); }
|
||||
static inline float4 ds4_hc_mix_widen(ushort4 v) { return ds4_hc_bf16x4_to_f32x4(v); }
|
||||
/* Share normalization and HC finalization, but preserve each weight type's
|
||||
* reference projection: F16 uses eight SIMDgroups per row, BF16 uses one. */
|
||||
template <typename W4, bool BF16 = false>
|
||||
static inline void ds4_hc_rms_norm_mix_cluster2_pre_norm_body(
|
||||
constant ds4_metal_args_hc_norm_mix & args,
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & split_args,
|
||||
@@ -1375,62 +1393,102 @@ static inline void ds4_hc_rms_norm_mix_cluster2_pre_norm_body(
|
||||
const float mean = total/(float)args.n;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
// Two independent eight-simdgroup clusters reproduce two original
|
||||
// NR0=2 matvec threadgroups inside this 512-thread threadgroup.
|
||||
const short cluster = sgitg / NSG_CLUSTER;
|
||||
const short local_sg = sgitg - cluster*NSG_CLUSTER;
|
||||
const int nb = args.n/NB;
|
||||
const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0;
|
||||
|
||||
device const half4 *ax4[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
ax4[row] = (device const half4 *)
|
||||
(weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half));
|
||||
}
|
||||
|
||||
float sumf_mv[NR0] = { 0.f };
|
||||
const short ix = tiisg/(NW/NF);
|
||||
const short il = tiisg%(NW/NF);
|
||||
const int ib0 = local_sg*NF + ix;
|
||||
for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) {
|
||||
float4 yl4[NF4];
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale;
|
||||
}
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
device const half4 *xb4 = ax4[row] + (ib*NB + il*NF)/4;
|
||||
float sumq = 0.f;
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
sumq += dot(float4(xb4[i]), yl4[i]);
|
||||
device volatile float *mixes_f32 = (device volatile float *)dst;
|
||||
if (BF16) {
|
||||
// GLM's scalar BF16 matvec owns a row in one SIMDgroup. Keep its
|
||||
// lane-strided operands, eight ordered FMAs, and single SIMD sum.
|
||||
if (sgitg < 4u) {
|
||||
const uint row = tgpig.x * 4u + sgitg;
|
||||
device const ushort *w = (device const ushort *)weight + (ulong)row * n;
|
||||
device const float *xr = (device const float *)x;
|
||||
float sum = 0.0f;
|
||||
for (uint k = tiisg; k < n; k += 256u) {
|
||||
const ushort w0 = w[k];
|
||||
const ushort w1 = w[k + 32u];
|
||||
const ushort w2 = w[k + 64u];
|
||||
const ushort w3 = w[k + 96u];
|
||||
const ushort w4 = w[k + 128u];
|
||||
const ushort w5 = w[k + 160u];
|
||||
const ushort w6 = w[k + 192u];
|
||||
const ushort w7 = w[k + 224u];
|
||||
// The unfused RMS kernel stores F32 before the projection.
|
||||
// Volatile preserves that rounding boundary under fast math.
|
||||
volatile float x0 = xr[k] * scale;
|
||||
volatile float x1 = xr[k + 32u] * scale;
|
||||
volatile float x2 = xr[k + 64u] * scale;
|
||||
volatile float x3 = xr[k + 96u] * scale;
|
||||
volatile float x4 = xr[k + 128u] * scale;
|
||||
volatile float x5 = xr[k + 160u] * scale;
|
||||
volatile float x6 = xr[k + 192u] * scale;
|
||||
volatile float x7 = xr[k + 224u] * scale;
|
||||
sum = fma(as_type<float>((uint)w0 << 16), x0, sum);
|
||||
sum = fma(as_type<float>((uint)w1 << 16), x1, sum);
|
||||
sum = fma(as_type<float>((uint)w2 << 16), x2, sum);
|
||||
sum = fma(as_type<float>((uint)w3 << 16), x3, sum);
|
||||
sum = fma(as_type<float>((uint)w4 << 16), x4, sum);
|
||||
sum = fma(as_type<float>((uint)w5 << 16), x5, sum);
|
||||
sum = fma(as_type<float>((uint)w6 << 16), x6, sum);
|
||||
sum = fma(as_type<float>((uint)w7 << 16), x7, sum);
|
||||
}
|
||||
sumf_mv[row] += sumq;
|
||||
sum = simd_sum(sum);
|
||||
if (tiisg == 0u) mixes_f32[row] = sum;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Two independent eight-simdgroup clusters reproduce two original
|
||||
// NR0=2 matvec threadgroups inside this 512-thread threadgroup.
|
||||
const short cluster = sgitg / NSG_CLUSTER;
|
||||
const short local_sg = sgitg - cluster*NSG_CLUSTER;
|
||||
const int nb = args.n/NB;
|
||||
const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0;
|
||||
|
||||
threadgroup float *cluster_shmem[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
cluster_shmem[row] = mv_shmem +
|
||||
((uint)cluster*NR0 + row)*NW;
|
||||
if (local_sg == 0) {
|
||||
cluster_shmem[row][tiisg] = 0.0f;
|
||||
}
|
||||
sumf_mv[row] = simd_sum(sumf_mv[row]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
cluster_shmem[row][local_sg] = sumf_mv[row];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
device volatile float *mixes_f32 =
|
||||
(device volatile float *)dst;
|
||||
if (local_sg == 0) {
|
||||
device const W4 *ax4[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
const float tot = simd_sum(cluster_shmem[row][tiisg]);
|
||||
if (tiisg == 0 && r0 + row < args.out_dim) {
|
||||
mixes_f32[r0 + row] = tot;
|
||||
ax4[row] = (device const W4 *)
|
||||
(weight + (uint64_t)(r0 + row)*(uint64_t)n*(sizeof(W4)/4));
|
||||
}
|
||||
|
||||
float sumf_mv[NR0] = { 0.f };
|
||||
const short ix = tiisg/(NW/NF);
|
||||
const short il = tiisg%(NW/NF);
|
||||
const int ib0 = local_sg*NF + ix;
|
||||
for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) {
|
||||
float4 yl4[NF4];
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale;
|
||||
}
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
device const W4 *xb4 = ax4[row] + (ib*NB + il*NF)/4;
|
||||
float sumq = 0.f;
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
sumq += dot(ds4_hc_mix_widen(xb4[i]), yl4[i]);
|
||||
}
|
||||
sumf_mv[row] += sumq;
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup float *cluster_shmem[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
cluster_shmem[row] = mv_shmem +
|
||||
((uint)cluster*NR0 + row)*NW;
|
||||
if (local_sg == 0) {
|
||||
cluster_shmem[row][tiisg] = 0.0f;
|
||||
}
|
||||
sumf_mv[row] = simd_sum(sumf_mv[row]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
cluster_shmem[row][local_sg] = sumf_mv[row];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (local_sg == 0) {
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
const float tot = simd_sum(cluster_shmem[row][tiisg]);
|
||||
if (tiisg == 0 && r0 + row < args.out_dim) {
|
||||
mixes_f32[r0 + row] = tot;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1560,7 +1618,7 @@ kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm(
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
||||
ds4_hc_rms_norm_mix_cluster2_pre_norm_body(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg);
|
||||
ds4_hc_rms_norm_mix_cluster2_pre_norm_body<half4>(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg);
|
||||
}
|
||||
|
||||
/* Decode-time fusion of the HC post/expand that follows a TP combine with the
|
||||
@@ -1631,5 +1689,28 @@ kernel void kernel_dsv4_hc_expand4_rms_norm_mix_f16_cluster2_pre_norm(
|
||||
thread_scope_device);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
|
||||
ds4_hc_rms_norm_mix_cluster2_pre_norm_body(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg);
|
||||
ds4_hc_rms_norm_mix_cluster2_pre_norm_body<half4>(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg);
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_rms_norm_mix_bf16_cluster2_pre_norm(
|
||||
constant ds4_metal_args_hc_norm_mix & args,
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & split_args,
|
||||
device const char * x,
|
||||
device const char * weight,
|
||||
device char * dst,
|
||||
device const float * hc_scale,
|
||||
device const float * hc_base,
|
||||
device char * split,
|
||||
device char * collapse_dst,
|
||||
device const char * norm_weight,
|
||||
device char * norm_dst,
|
||||
device atomic_uint * completion,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
||||
ds4_hc_rms_norm_mix_cluster2_pre_norm_body<ushort4, true>(
|
||||
args, split_args, x, weight, dst, hc_scale, hc_base, split,
|
||||
collapse_dst, norm_weight, norm_dst, completion, shmem,
|
||||
tgpig, tiisg, sgitg);
|
||||
}
|
||||
+790
-2
@@ -365,6 +365,19 @@ struct ds4_metal_args_glm_attention_indexed_decode_split {
|
||||
uint32_t value_type;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_glm_attention_indexed_decode_exact {
|
||||
uint32_t n_selected;
|
||||
uint32_t cache_cap;
|
||||
uint32_t n_head;
|
||||
uint32_t kv_lora_dim;
|
||||
uint32_t value_dim;
|
||||
uint32_t value_row_bytes;
|
||||
uint32_t value_type;
|
||||
uint32_t stage_rows; /* rows staged per threadgroup by the score kernel */
|
||||
uint32_t heads_per_group; /* heads per threadgroup in the lora kernel */
|
||||
float scale;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_glm_attention_indexed_batch {
|
||||
uint32_t n_tokens;
|
||||
uint32_t n_selected;
|
||||
@@ -2356,6 +2369,32 @@ kernel void kernel_glm_qk_lowrank_q8_0(
|
||||
}
|
||||
}
|
||||
|
||||
// Tile independent output columns; each dot retains the reference operands
|
||||
// and accumulation order. GLM-5.3 dispatches one SIMDgroup per column tile.
|
||||
kernel void kernel_glm_qk_lowrank_q8_0_tiled(
|
||||
constant ds4_metal_args_glm_qk_lowrank & args,
|
||||
device const char *weight,
|
||||
device const char *q,
|
||||
device char *qk_low,
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg_u [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const uint head = tgpig.x;
|
||||
if (head >= args.n_head) return;
|
||||
const uint nth = ntg_u.x;
|
||||
device const float *qh =
|
||||
(device const float *)(q + (uint64_t)head * args.qk_dim * sizeof(float));
|
||||
device float *out =
|
||||
(device float *)(qk_low + (uint64_t)head * args.kv_lora_dim * sizeof(float));
|
||||
|
||||
const uint j = tgpig.y * nth + tid;
|
||||
if (j < args.kv_lora_dim) {
|
||||
device const char *row =
|
||||
weight + ((uint64_t)head * args.kv_lora_dim + j) * args.row_bytes;
|
||||
out[j] = glm_quant_dot_row_dev_f32(args.weight_type, row, qh, args.qk_nope);
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_glm_qk_lowrank_q8_0_glm52(
|
||||
constant ds4_metal_args_glm_qk_lowrank & args,
|
||||
device const char *weight,
|
||||
@@ -2609,6 +2648,123 @@ kernel void kernel_glm_qk_lowrank_q8_0_batch_glm52_t4(
|
||||
}
|
||||
}
|
||||
|
||||
/* GLM 5.3 prefill qk-low, one threadgroup per (head, TT tokens).
|
||||
*
|
||||
* kernel_glm_qk_lowrank_q8_0_batch gives every (head, token) pair its own
|
||||
* threadgroup, so each of the 512 Q8_0 rows of a head's K_b slice is re-read
|
||||
* once per token -- 17.8 GB of weight traffic per layer for 34 GFLOP -- and
|
||||
* each thread runs a single 256-long dependent FMA chain.
|
||||
*
|
||||
* Here a thread owns the same two output rows for TT consecutive tokens: it
|
||||
* walks its 272-byte row once per TT tokens and keeps TT independent
|
||||
* accumulator chains. Which threadgroup computes which outputs changes;
|
||||
* the arithmetic does not. Every output still accumulates
|
||||
* `acc += d * (float)qs[qi] * x[base + qi]` over ascending blocks and
|
||||
* ascending qi, written verbatim, so the result is bit-identical to the
|
||||
* reference kernel above.
|
||||
*/
|
||||
template <uint TT>
|
||||
static inline void glm_qk_lowrank_q8_0_batch_tokens_impl(
|
||||
constant ds4_metal_args_glm_qk_lowrank_batch & args,
|
||||
device const char *weight,
|
||||
device const char *q,
|
||||
device char *qk_low,
|
||||
uint tid,
|
||||
uint nth,
|
||||
uint3 tgpig) {
|
||||
constexpr uint kv_lora_dim = 512u;
|
||||
constexpr uint qk_nope = 256u;
|
||||
constexpr uint qk_dim = 256u;
|
||||
constexpr uint row_bytes = 272u;
|
||||
constexpr uint n_blocks = qk_nope >> 5;
|
||||
constexpr uint NR = 2u;
|
||||
|
||||
if (args.kv_lora_dim != kv_lora_dim ||
|
||||
args.qk_nope != qk_nope ||
|
||||
args.qk_dim != qk_dim ||
|
||||
args.row_bytes != row_bytes ||
|
||||
args.weight_type != DS4_METAL_GGUF_Q8_0 ||
|
||||
args.n_tokens == 0u ||
|
||||
nth * NR != kv_lora_dim) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint head = tgpig.x + args.head_base;
|
||||
const uint token0 = tgpig.y * TT;
|
||||
if (head >= args.n_head || token0 >= args.n_tokens) return;
|
||||
|
||||
const ulong q_token_stride = (ulong)args.n_head * qk_dim;
|
||||
const ulong low_token_stride = (ulong)args.n_head * kv_lora_dim;
|
||||
device const float *xbase = (device const float *)q;
|
||||
|
||||
/* A partial tail tile clamps to the last real token instead of branching:
|
||||
* the clamped columns are read and accumulated, then never stored. */
|
||||
ulong xoff[TT];
|
||||
FOR_UNROLL (uint t = 0; t < TT; t++) {
|
||||
const uint token = min(token0 + t, args.n_tokens - 1u);
|
||||
xoff[t] = token * q_token_stride + head * qk_dim;
|
||||
}
|
||||
|
||||
const uint j0 = tid;
|
||||
const uint j1 = tid + nth;
|
||||
device const char *row0 =
|
||||
weight + ((uint64_t)head * kv_lora_dim + j0) * row_bytes;
|
||||
device const char *row1 =
|
||||
weight + ((uint64_t)head * kv_lora_dim + j1) * row_bytes;
|
||||
|
||||
float acc0[TT];
|
||||
float acc1[TT];
|
||||
FOR_UNROLL (uint t = 0; t < TT; t++) {
|
||||
acc0[t] = 0.0f;
|
||||
acc1[t] = 0.0f;
|
||||
}
|
||||
|
||||
for (uint block = 0; block < n_blocks; block++) {
|
||||
device const char *block0 = row0 + (uint64_t)block * 34u;
|
||||
device const char *block1 = row1 + (uint64_t)block * 34u;
|
||||
const float d0 = (float)(*((device const half *)block0));
|
||||
const float d1 = (float)(*((device const half *)block1));
|
||||
device const int8_t *qs0 = (device const int8_t *)(block0 + 2u);
|
||||
device const int8_t *qs1 = (device const int8_t *)(block1 + 2u);
|
||||
const uint base = block << 5;
|
||||
for (uint qi = 0; qi < 32u; qi++) {
|
||||
const uint col = base + qi;
|
||||
FOR_UNROLL (uint t = 0; t < TT; t++) {
|
||||
acc0[t] += d0 * (float)qs0[qi] * xbase[xoff[t] + col];
|
||||
acc1[t] += d1 * (float)qs1[qi] * xbase[xoff[t] + col];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device float *outbase = (device float *)qk_low;
|
||||
FOR_UNROLL (uint t = 0; t < TT; t++) {
|
||||
const uint token = token0 + t;
|
||||
if (token < args.n_tokens) {
|
||||
device float *out =
|
||||
outbase + token * low_token_stride + head * kv_lora_dim;
|
||||
out[j0] = acc0[t];
|
||||
out[j1] = acc1[t];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(TT) \
|
||||
kernel void kernel_glm_qk_lowrank_q8_0_batch_t##TT( \
|
||||
constant ds4_metal_args_glm_qk_lowrank_batch & args, \
|
||||
device const char *weight, \
|
||||
device const char *q, \
|
||||
device char *qk_low, \
|
||||
uint tid [[thread_index_in_threadgroup]], \
|
||||
ushort3 ntg_u [[threads_per_threadgroup]], \
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) { \
|
||||
glm_qk_lowrank_q8_0_batch_tokens_impl<TT>( \
|
||||
args, weight, q, qk_low, tid, ntg_u.x, tgpig); \
|
||||
}
|
||||
|
||||
DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(4)
|
||||
DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(8)
|
||||
DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(16)
|
||||
|
||||
kernel void kernel_glm_value_project_q8_0(
|
||||
constant ds4_metal_args_glm_qk_lowrank & args,
|
||||
device const char *weight,
|
||||
@@ -2868,7 +3024,12 @@ kernel void kernel_glm_attention_indexed_decode_split_group8_partial_impl(
|
||||
if (args.n_selected == 0u ||
|
||||
args.cache_f16 == 0u ||
|
||||
args.kv_lora_dim != 512u ||
|
||||
args.qk_rope != 64u ||
|
||||
/* GLM 5.3 has no RoPE tail (n_rot = 0). Everything rope here is
|
||||
* driven by rope_vecs = qk_rope >> 2, so at 0 the staging loop runs no
|
||||
* iterations, rope_shared is never touched and the per-lane rope dot
|
||||
* is skipped -- the kernel is already correct for that case and only
|
||||
* this guard kept it out. */
|
||||
(args.qk_rope != 64u && args.qk_rope != 0u) ||
|
||||
args.block_rows == 0u ||
|
||||
block >= args.n_blocks) {
|
||||
return;
|
||||
@@ -3178,6 +3339,260 @@ kernel void kernel_glm_attention_indexed_decode_split_group8_reduce16(
|
||||
tid, ntg_u, tgpig);
|
||||
}
|
||||
|
||||
/*
|
||||
* kernel_glm_attention_indexed_decode, in phases, with its arithmetic kept
|
||||
* operation for operation.
|
||||
*
|
||||
* The generic kernel below runs one threadgroup per head and has every head
|
||||
* walk every selected row twice, so 64 heads re-read each cache row 128 times
|
||||
* and only 64 threadgroups exist to do it. The four kernels here compute the
|
||||
* same thing in the same order -- one thread per (head, row) for the
|
||||
* sequential 512-term score, the generic's 256-thread partition and
|
||||
* 128/64/../1 tree for the softmax denominator, one thread per column pair
|
||||
* walking rows 0..n-1 for the weighted sum, and the same quantised value row
|
||||
* dot -- but stage each cache row once per threadgroup for every head and
|
||||
* spread the work over hundreds of threadgroups. Every floating-point
|
||||
* operation, operand and ordering is the generic kernel's, so the output is
|
||||
* bit-identical to it (tests/test_glm53_kda asserts this at tolerance 0);
|
||||
* the speed comes from row sharing and parallelism alone. Intermediates live
|
||||
* in device buffers instead of threadgroup memory: scores[head][s], which the
|
||||
* weights kernel turns into softmax weights in place, denom[head] and
|
||||
* lora[head][kv_lora_dim]. f16 compact cache and no RoPE tail only, which is
|
||||
* GLM 5.3.
|
||||
*
|
||||
* Score kernel. Grid: ceil(n_selected / stage_rows) threadgroups of
|
||||
* n_head * stage_rows threads. Thread t scores row t % stage_rows for head
|
||||
* t / stage_rows, so a simdgroup holds two heads over the same rows and the
|
||||
* staged half4 it reads are shared lane pairs. Rows are staged at an odd
|
||||
* half4 stride so consecutive rows fall in different banks.
|
||||
*/
|
||||
kernel void kernel_glm_attention_indexed_decode_exact_scores(
|
||||
constant ds4_metal_args_glm_attention_indexed_decode_exact & args,
|
||||
device const char *qk_low,
|
||||
device const char *kv_lora_cache,
|
||||
device const uint32_t *selected,
|
||||
device float *scores,
|
||||
threadgroup half4 *kv_shared [[threadgroup(0)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const uint stage_rows = args.stage_rows;
|
||||
const uint kv_vecs = args.kv_lora_dim >> 2;
|
||||
const uint row_stride = kv_vecs | 1u;
|
||||
const uint s0 = tgpig.x * stage_rows;
|
||||
if (s0 >= args.n_selected || stage_rows == 0u) return;
|
||||
const uint rows = min(stage_rows, args.n_selected - s0);
|
||||
const uint nthreads = args.n_head * stage_rows;
|
||||
for (uint off = tid; off < rows * kv_vecs; off += nthreads) {
|
||||
const uint rr = off / kv_vecs;
|
||||
const uint vv = off - rr * kv_vecs;
|
||||
const uint row = selected[s0 + rr];
|
||||
kv_shared[rr * row_stride + vv] = row < args.cache_cap
|
||||
? ((device const half4 *)kv_lora_cache)[(uint64_t)row * kv_vecs + vv]
|
||||
: half4(half(0.0f));
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const uint r = tid % stage_rows;
|
||||
const uint head = tid / stage_rows;
|
||||
if (r >= rows || head >= args.n_head) return;
|
||||
const uint s = s0 + r;
|
||||
const uint row = selected[s];
|
||||
device const float *low =
|
||||
(device const float *)(qk_low + (uint64_t)head * args.kv_lora_dim * sizeof(float));
|
||||
threadgroup const half4 *kvrow = kv_shared + r * row_stride;
|
||||
float score = -INFINITY;
|
||||
if (row < args.cache_cap) {
|
||||
float dotv = 0.0f;
|
||||
uint j = 0;
|
||||
for (; j + 3u < args.kv_lora_dim; j += 4u) {
|
||||
threadgroup const half4 *kv4 = kvrow + (j >> 2);
|
||||
device const float4 *low4 =
|
||||
(device const float4 *)(low + j);
|
||||
const float4 kv = (float4)(*kv4);
|
||||
const float4 qv = *low4;
|
||||
dotv += qv.x * kv.x + qv.y * kv.y +
|
||||
qv.z * kv.z + qv.w * kv.w;
|
||||
}
|
||||
if (j < args.kv_lora_dim) {
|
||||
for (; j < args.kv_lora_dim; j++) {
|
||||
const float kv = (float)((threadgroup const half *)kvrow)[j];
|
||||
dotv += low[j] * kv;
|
||||
}
|
||||
}
|
||||
score = dotv * args.scale;
|
||||
}
|
||||
scores[(uint64_t)head * args.n_selected + s] = score;
|
||||
}
|
||||
|
||||
/* Weights kernel. Grid: n_head threadgroups of 256 threads -- the generic
|
||||
* kernel's threadgroup -- so the per-thread row partition and the reduction
|
||||
* tree produce its max and denominator. Scores become softmax weights in
|
||||
* place. */
|
||||
kernel void kernel_glm_attention_indexed_decode_exact_weights(
|
||||
constant ds4_metal_args_glm_attention_indexed_decode_exact & args,
|
||||
device float *scores,
|
||||
device float *denom,
|
||||
threadgroup float *red [[threadgroup(0)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg_u [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const uint head = tgpig.x;
|
||||
if (head >= args.n_head || args.n_selected == 0u) return;
|
||||
const uint nth = ntg_u.x;
|
||||
device float *sc = scores + (uint64_t)head * args.n_selected;
|
||||
|
||||
float local_max = -INFINITY;
|
||||
for (uint s = tid; s < args.n_selected; s += nth) {
|
||||
local_max = max(local_max, sc[s]);
|
||||
}
|
||||
red[tid] = local_max;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint step = nth >> 1; step > 0; step >>= 1) {
|
||||
if (tid < step) red[tid] = max(red[tid], red[tid + step]);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
const float max_score = red[0];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float local_sum = 0.0f;
|
||||
for (uint s = tid; s < args.n_selected; s += nth) {
|
||||
const float w = exp(sc[s] - max_score);
|
||||
sc[s] = w;
|
||||
local_sum += w;
|
||||
}
|
||||
red[tid] = local_sum;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint step = nth >> 1; step > 0; step >>= 1) {
|
||||
if (tid < step) red[tid] += red[tid + step];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
if (tid == 0u) denom[head] = max(red[0], 1.0e-20f);
|
||||
}
|
||||
|
||||
/* Lora kernel. Grid: ceil(n_head / heads_per_group) x (kv_lora_dim / 64)
|
||||
* threadgroups of heads_per_group * 32 threads. A simdgroup is one head over
|
||||
* 64 consecutive columns, two per lane as in the generic kernel. The rows
|
||||
* are walked in stages of 32: every thread fetches 16 bytes of the stage's
|
||||
* 128-byte row slices and one weight, three stages ahead of use, and parks
|
||||
* them in double-buffered threadgroup memory, so the scattered row reads are
|
||||
* in flight while the fma chains run. Those chains are the generic kernel's,
|
||||
* row after row in selection order. */
|
||||
#define DS4_GLM_EXACT_LORA_STAGE 32u
|
||||
#define DS4_GLM_EXACT_LORA_AHEAD 3u
|
||||
kernel void kernel_glm_attention_indexed_decode_exact_lora(
|
||||
constant ds4_metal_args_glm_attention_indexed_decode_exact & args,
|
||||
device const char *kv_lora_cache,
|
||||
device const uint32_t *selected,
|
||||
device const float *weights,
|
||||
device const float *denom,
|
||||
device float *lora,
|
||||
threadgroup uchar *scratch [[threadgroup(0)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
constexpr uint stage = DS4_GLM_EXACT_LORA_STAGE;
|
||||
constexpr uint ahead = DS4_GLM_EXACT_LORA_AHEAD;
|
||||
const uint lane = tid & 31u;
|
||||
const uint sg = tid >> 5u;
|
||||
const uint head0 = tgpig.x * args.heads_per_group;
|
||||
const uint head = head0 + sg;
|
||||
const uint c0 = tgpig.y * 64u;
|
||||
const uint j0 = c0 + lane * 2u;
|
||||
const uint hpg = args.heads_per_group;
|
||||
const uint n = args.n_selected;
|
||||
const uint nstages = (n + stage - 1u) / stage;
|
||||
device const half *cache = (device const half *)kv_lora_cache;
|
||||
|
||||
/* Staging roles: thread t fetches row t / 8, 16-byte chunk t % 8 of the
|
||||
* kv slice, and the weight of head t / 32, row t % 32. */
|
||||
const uint kv_r = tid >> 3u;
|
||||
const uint kv_chunk = tid & 7u;
|
||||
const uint w_h = tid >> 5u;
|
||||
const uint w_r = tid & 31u;
|
||||
const uint kv_slice_bytes = stage * 64u * sizeof(half);
|
||||
const uint w_slice_bytes = hpg * stage * sizeof(float);
|
||||
threadgroup uchar *buf0 = scratch;
|
||||
threadgroup uchar *buf1 = scratch + kv_slice_bytes + w_slice_bytes;
|
||||
|
||||
uint4 kvq[ahead];
|
||||
float wq[ahead];
|
||||
#define DS4_GLM_EXACT_LORA_FETCH(k, slot) do { \
|
||||
const uint s0_ = (k) * stage; \
|
||||
const uint s_ = s0_ + kv_r; \
|
||||
const uint row_ = s_ < n ? selected[s_] : 0u; \
|
||||
if (s_ < n && row_ < args.cache_cap) { \
|
||||
kvq[slot] = *(device const uint4 *)(cache + (uint64_t)row_ * args.kv_lora_dim + c0 + kv_chunk * 8u); \
|
||||
} else { \
|
||||
kvq[slot] = uint4(0u); \
|
||||
} \
|
||||
const uint sw_ = s0_ + w_r; \
|
||||
const uint hh_ = head0 + w_h; \
|
||||
wq[slot] = (sw_ < n && hh_ < args.n_head) \
|
||||
? weights[(uint64_t)hh_ * n + sw_] : 0.0f; \
|
||||
} while (0)
|
||||
|
||||
for (uint i = 0; i < ahead; i++) {
|
||||
if (i < nstages) DS4_GLM_EXACT_LORA_FETCH(i, i);
|
||||
}
|
||||
float acc0 = 0.0f;
|
||||
float acc1 = 0.0f;
|
||||
for (uint k = 0; k < nstages; k++) {
|
||||
threadgroup uchar *buf = (k & 1u) ? buf1 : buf0;
|
||||
threadgroup uint4 *kv_sh = (threadgroup uint4 *)buf;
|
||||
threadgroup float *w_sh = (threadgroup float *)(buf + kv_slice_bytes);
|
||||
kv_sh[kv_r * 8u + kv_chunk] = kvq[0];
|
||||
w_sh[w_h * stage + w_r] = wq[0];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint i = 1; i < ahead; i++) {
|
||||
kvq[i - 1u] = kvq[i];
|
||||
wq[i - 1u] = wq[i];
|
||||
}
|
||||
if (k + ahead < nstages) DS4_GLM_EXACT_LORA_FETCH(k + ahead, ahead - 1u);
|
||||
if (head < args.n_head) {
|
||||
threadgroup const half2 *kv_rows = (threadgroup const half2 *)buf;
|
||||
threadgroup const float *w_rows = w_sh + sg * stage;
|
||||
for (uint r = 0; r < stage; r++) {
|
||||
const float2 kv = (float2)kv_rows[r * 32u + lane];
|
||||
const float w = w_rows[r];
|
||||
acc0 += w * kv.x;
|
||||
acc1 += w * kv.y;
|
||||
}
|
||||
}
|
||||
}
|
||||
#undef DS4_GLM_EXACT_LORA_FETCH
|
||||
if (head >= args.n_head) return;
|
||||
const float d = denom[head];
|
||||
device float *out = lora + (uint64_t)head * args.kv_lora_dim;
|
||||
out[j0] = acc0 / d;
|
||||
out[j0 + 1u] = acc1 / d;
|
||||
}
|
||||
|
||||
/* Value kernel. Grid: n_head x ceil(value_dim / threads) threadgroups; each
|
||||
* copies its head's lora vector into threadgroup memory and runs the generic
|
||||
* kernel's quantised row dot for one output element per thread. */
|
||||
kernel void kernel_glm_attention_indexed_decode_exact_value(
|
||||
constant ds4_metal_args_glm_attention_indexed_decode_exact & args,
|
||||
device const float *lora,
|
||||
device const char *value_weight,
|
||||
device char *heads,
|
||||
threadgroup float *lora_sum [[threadgroup(0)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg_u [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const uint head = tgpig.x;
|
||||
if (head >= args.n_head) return;
|
||||
const uint nth = ntg_u.x;
|
||||
device const float *src = lora + (uint64_t)head * args.kv_lora_dim;
|
||||
for (uint j = tid; j < args.kv_lora_dim; j += nth) lora_sum[j] = src[j];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint d = tgpig.y * nth + tid;
|
||||
if (d >= args.value_dim) return;
|
||||
device float *out =
|
||||
(device float *)(heads + (uint64_t)head * args.value_dim * sizeof(float));
|
||||
device const char *row =
|
||||
value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes;
|
||||
out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, lora_sum, args.kv_lora_dim);
|
||||
}
|
||||
|
||||
kernel void kernel_glm_attention_indexed_decode(
|
||||
constant ds4_metal_args_glm_attention_indexed_decode & args,
|
||||
device const char *q,
|
||||
@@ -3785,6 +4200,8 @@ kernel void kernel_glm_attention_indexed_batch_group2(
|
||||
}
|
||||
}
|
||||
|
||||
/* Original one-head reference from 8969dbb. Keep its expressions and
|
||||
* specializations independent of the wider-head implementation below. */
|
||||
template <bool assume_valid_rows, bool assume_valid_heads>
|
||||
kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl(
|
||||
constant ds4_metal_args_glm_attention_indexed_batch & args,
|
||||
@@ -3986,6 +4403,233 @@ template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_f
|
||||
kernel glm_attention_indexed_batch_lora_group8_vec_t
|
||||
kernel_glm_attention_indexed_batch_lora_group8_vec_impl<true, true>;
|
||||
|
||||
/* Indexed prefill attention over the compact KV cache.
|
||||
*
|
||||
* A threadgroup is 8 simdgroups; heads_per_sg is how many heads one simdgroup
|
||||
* carries, so the threadgroup covers 8 * heads_per_sg heads and every token
|
||||
* needs 64 / (8 * heads_per_sg) staging passes over its selected rows. At
|
||||
* heads_per_sg 1 each of the 8 head groups re-stages all 2051 selected rows
|
||||
* (2 MB per token), which is 34 GB per layer at a 2048-token chunk; carrying
|
||||
* two heads per simdgroup halves that and each staged row block feeds both.
|
||||
*
|
||||
* Nothing a head computes depends on heads_per_sg: the row order, the four
|
||||
* dot(float4) terms, the simd_sum tree and the online-softmax update are the
|
||||
* same expressions in the same order, so every head's output is bit-identical
|
||||
* across the instantiations. */
|
||||
template <bool assume_valid_rows, bool assume_valid_heads, uint heads_per_sg>
|
||||
kernel void kernel_glm_attention_indexed_batch_lora_heads_impl(
|
||||
constant ds4_metal_args_glm_attention_indexed_batch & args,
|
||||
device const char *q,
|
||||
device const char *qk_low,
|
||||
device const char *kv_lora_cache,
|
||||
device const char *k_rope_cache,
|
||||
device const uint32_t *selected,
|
||||
device char *lora_out,
|
||||
threadgroup half4 *scratch [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tid_u [[thread_index_in_threadgroup]],
|
||||
ushort lane_u [[thread_index_in_simdgroup]],
|
||||
ushort sg_u [[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr uint group_heads = 8u;
|
||||
constexpr uint stage_rows = 16u;
|
||||
const uint token = tgpig.y;
|
||||
const uint tid = (uint)tid_u;
|
||||
const uint lane = (uint)lane_u;
|
||||
const uint head_in_group = (uint)sg_u;
|
||||
const uint head0 =
|
||||
(tgpig.x * group_heads + head_in_group) * heads_per_sg + args.head_base;
|
||||
if (token >= args.n_tokens ||
|
||||
args.n_selected == 0u ||
|
||||
args.cache_f16 == 0u ||
|
||||
args.kv_lora_dim != 512u ||
|
||||
(args.qk_rope != 0u && args.qk_rope != 64u)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint kv_vecs = args.kv_lora_dim >> 2;
|
||||
const uint rope_vecs = args.qk_rope >> 2;
|
||||
const uint qk_dim = args.qk_nope + args.qk_rope;
|
||||
const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float);
|
||||
const uint64_t low_token_stride =
|
||||
(uint64_t)args.n_head * args.kv_lora_dim * sizeof(float);
|
||||
|
||||
threadgroup half4 *kv_shared = scratch;
|
||||
threadgroup float4 *rope_shared =
|
||||
(threadgroup float4 *)(kv_shared + stage_rows * kv_vecs);
|
||||
|
||||
bool valid_head[heads_per_sg];
|
||||
float4 low[heads_per_sg][4];
|
||||
float4 qrope[heads_per_sg];
|
||||
FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) {
|
||||
const uint head = head0 + k;
|
||||
valid_head[k] = assume_valid_heads || head < args.n_head;
|
||||
const uint safe_head = valid_head[k] ? head : 0u;
|
||||
device const float *qh =
|
||||
(device const float *)(q +
|
||||
(uint64_t)token * q_token_stride +
|
||||
(uint64_t)safe_head * qk_dim * sizeof(float));
|
||||
device const float4 *low4 =
|
||||
(device const float4 *)(qk_low +
|
||||
(uint64_t)token * low_token_stride +
|
||||
(uint64_t)safe_head * args.kv_lora_dim * sizeof(float));
|
||||
low[k][0] = 0.0f;
|
||||
low[k][1] = 0.0f;
|
||||
low[k][2] = 0.0f;
|
||||
low[k][3] = 0.0f;
|
||||
qrope[k] = 0.0f;
|
||||
if (valid_head[k]) {
|
||||
low[k][0] = low4[lane + 0u];
|
||||
low[k][1] = low4[lane + 32u];
|
||||
low[k][2] = low4[lane + 64u];
|
||||
low[k][3] = low4[lane + 96u];
|
||||
if (lane < rope_vecs) {
|
||||
qrope[k] = *((device const float4 *)(qh + args.qk_nope + lane * 4u));
|
||||
}
|
||||
}
|
||||
}
|
||||
device const uint32_t *token_selected =
|
||||
selected + (uint64_t)token * args.n_selected;
|
||||
|
||||
float corr_dims[2] = {0.0f, 0.0f};
|
||||
if (args.qk_rope != 0u && args.ext_factor != 0.0f) {
|
||||
glm_rope_yarn_corr_dims((int)args.qk_rope,
|
||||
(int)args.n_ctx_orig,
|
||||
args.freq_base,
|
||||
args.beta_fast,
|
||||
args.beta_slow,
|
||||
corr_dims);
|
||||
}
|
||||
|
||||
float M[heads_per_sg];
|
||||
float S[heads_per_sg];
|
||||
float4 o[heads_per_sg][4];
|
||||
FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) {
|
||||
M[k] = -FLT_MAX / 2.0f;
|
||||
S[k] = 0.0f;
|
||||
o[k][0] = 0.0f;
|
||||
o[k][1] = 0.0f;
|
||||
o[k][2] = 0.0f;
|
||||
o[k][3] = 0.0f;
|
||||
}
|
||||
|
||||
for (uint base = 0u; base < args.n_selected; base += stage_rows) {
|
||||
const uint rows = min(stage_rows, args.n_selected - base);
|
||||
for (uint off = tid; off < rows * kv_vecs; off += 256u) {
|
||||
const uint rr = off / kv_vecs;
|
||||
const uint vv = off - rr * kv_vecs;
|
||||
const uint row = token_selected[base + rr];
|
||||
const bool valid_row = assume_valid_rows || row < args.cache_cap;
|
||||
if (valid_row) {
|
||||
device const half4 *src =
|
||||
(device const half4 *)((device const half *)kv_lora_cache +
|
||||
(uint64_t)row * args.kv_lora_dim);
|
||||
kv_shared[off] = src[vv];
|
||||
} else {
|
||||
kv_shared[off] = half4(half(0.0f));
|
||||
}
|
||||
}
|
||||
for (uint off = tid; off < rows * rope_vecs; off += 256u) {
|
||||
const uint rr = off / rope_vecs;
|
||||
const uint vv = off - rr * rope_vecs;
|
||||
const uint r = vv * 4u;
|
||||
const uint row = token_selected[base + rr];
|
||||
const bool valid_row = assume_valid_rows || row < args.cache_cap;
|
||||
if (valid_row) {
|
||||
const uint64_t rope_base = (uint64_t)row * args.qk_rope;
|
||||
const float2 y0 =
|
||||
glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache,
|
||||
rope_base,
|
||||
r,
|
||||
row,
|
||||
args.qk_rope,
|
||||
args.freq_base,
|
||||
args.freq_scale,
|
||||
args.ext_factor,
|
||||
args.attn_factor,
|
||||
corr_dims[0],
|
||||
corr_dims[1]);
|
||||
const float2 y1 =
|
||||
glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache,
|
||||
rope_base,
|
||||
r + 2u,
|
||||
row,
|
||||
args.qk_rope,
|
||||
args.freq_base,
|
||||
args.freq_scale,
|
||||
args.ext_factor,
|
||||
args.attn_factor,
|
||||
corr_dims[0],
|
||||
corr_dims[1]);
|
||||
rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y);
|
||||
} else {
|
||||
rope_shared[off] = float4(0.0f);
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint rr = 0u; rr < rows; rr++) {
|
||||
const uint row = token_selected[base + rr];
|
||||
const bool valid_row = assume_valid_rows || row < args.cache_cap;
|
||||
threadgroup const half4 *kv_row = kv_shared + rr * kv_vecs;
|
||||
threadgroup const float4 *rope_row = rope_shared + rr * rope_vecs;
|
||||
/* Verbatim from the one-head kernel, including the repeated
|
||||
* `(float4)kv_row[...]` in the score and in the output update:
|
||||
* naming the converted row once is the same value but lets the
|
||||
* compiler contract the update the other way round, which does
|
||||
* not round the same. */
|
||||
FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) {
|
||||
float partial = 0.0f;
|
||||
if (valid_head[k] && valid_row) {
|
||||
partial += dot(low[k][0], (float4)kv_row[lane + 0u]);
|
||||
partial += dot(low[k][1], (float4)kv_row[lane + 32u]);
|
||||
partial += dot(low[k][2], (float4)kv_row[lane + 64u]);
|
||||
partial += dot(low[k][3], (float4)kv_row[lane + 96u]);
|
||||
if (lane < rope_vecs) {
|
||||
partial += dot(qrope[k], rope_row[lane]);
|
||||
}
|
||||
}
|
||||
const float sum = simd_sum(partial);
|
||||
const float score =
|
||||
(valid_head[k] && valid_row) ? sum * args.scale : -FLT_MAX / 2.0f;
|
||||
if (valid_head[k] && valid_row) {
|
||||
const float new_m = max(M[k], score);
|
||||
const float old_scale = exp(M[k] - new_m);
|
||||
const float row_scale = exp(score - new_m);
|
||||
o[k][0] = o[k][0] * old_scale + (float4)kv_row[lane + 0u] * row_scale;
|
||||
o[k][1] = o[k][1] * old_scale + (float4)kv_row[lane + 32u] * row_scale;
|
||||
o[k][2] = o[k][2] * old_scale + (float4)kv_row[lane + 64u] * row_scale;
|
||||
o[k][3] = o[k][3] * old_scale + (float4)kv_row[lane + 96u] * row_scale;
|
||||
S[k] = S[k] * old_scale + row_scale;
|
||||
M[k] = new_m;
|
||||
}
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) {
|
||||
if (valid_head[k]) {
|
||||
const float inv_s = S[k] > 0.0f ? 1.0f / S[k] : 0.0f;
|
||||
device float4 *out4 =
|
||||
(device float4 *)(lora_out +
|
||||
((uint64_t)token * args.n_head + head0 + k) *
|
||||
args.kv_lora_dim * sizeof(float));
|
||||
out4[lane + 0u] = o[k][0] * inv_s;
|
||||
out4[lane + 32u] = o[k][1] * inv_s;
|
||||
out4[lane + 64u] = o[k][2] * inv_s;
|
||||
out4[lane + 96u] = o[k][3] * inv_s;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template [[host_name("kernel_glm_attention_indexed_batch_lora_group16_vec_fullheads")]]
|
||||
kernel glm_attention_indexed_batch_lora_group8_vec_t
|
||||
kernel_glm_attention_indexed_batch_lora_heads_impl<false, true, 2u>;
|
||||
|
||||
template [[host_name("kernel_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads")]]
|
||||
kernel glm_attention_indexed_batch_lora_group8_vec_t
|
||||
kernel_glm_attention_indexed_batch_lora_heads_impl<true, true, 2u>;
|
||||
|
||||
template <bool assume_valid_heads>
|
||||
kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
|
||||
constant ds4_metal_args_glm_attention_indexed_batch & args,
|
||||
@@ -4819,7 +5463,8 @@ kernel void kernel_glm_router_select_one(
|
||||
if (tid < k_used) {
|
||||
token_selected[tid] = idx[tid];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// Normalization reads probabilities and selected IDs written by other threads.
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device);
|
||||
|
||||
if (tid < k_used) {
|
||||
float sum = 0.0f;
|
||||
@@ -4831,6 +5476,149 @@ kernel void kernel_glm_router_select_one(
|
||||
}
|
||||
}
|
||||
|
||||
// GLM-5.3's 288 experts fit in nine registers per lane of one SIMDgroup.
|
||||
static inline void ds4_glm_router_select_top8_impl(
|
||||
constant ds4_metal_args_glm_router_select_one & args,
|
||||
device const float *logits,
|
||||
device const float *bias,
|
||||
device int32_t *selected,
|
||||
device float *weights,
|
||||
device float *probs,
|
||||
threadgroup float *scratch,
|
||||
uint token,
|
||||
uint lane) {
|
||||
threadgroup float *scores = scratch;
|
||||
threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 512u);
|
||||
device const float *token_logits = logits + (ulong)token * args.n_expert;
|
||||
device int32_t *token_selected = selected + (ulong)token * args.n_expert_used;
|
||||
device float *token_weights = weights + (ulong)token * args.n_expert_used;
|
||||
device float *token_probs = probs + (ulong)token * args.n_expert;
|
||||
float values[9];
|
||||
bool has_nan = false;
|
||||
for (uint i = 0u; i < 9u; ++i) {
|
||||
const uint expert = lane + 32u * i;
|
||||
const float p = ds4_glm_router_sigmoid(token_logits[expert]);
|
||||
token_probs[expert] = p;
|
||||
values[i] = p + bias[expert];
|
||||
scores[expert] = values[i];
|
||||
has_nan |= (as_type<uint>(values[i]) & 0x7fffffffu) > 0x7f800000u;
|
||||
}
|
||||
if (simd_any(has_nan)) {
|
||||
// Emulate the same 512-wire sorting network, including its NaN behavior.
|
||||
for (uint tid = lane; tid < 512u; tid += 32u) {
|
||||
if (tid >= 288u) scores[tid] = -INFINITY;
|
||||
idx[tid] = (int32_t)tid;
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint k = 2u; k <= 512u; k <<= 1u) {
|
||||
for (uint j = k >> 1u; j > 0u; j >>= 1u) {
|
||||
for (uint tid = lane; tid < 512u; tid += 32u) {
|
||||
const uint other = tid ^ j;
|
||||
if (other > tid) {
|
||||
const int32_t a = idx[tid];
|
||||
const int32_t b = idx[other];
|
||||
const bool descending = (tid & k) == 0u;
|
||||
const bool swap = descending ? ds4_glm_router_better(scores, b, a)
|
||||
: ds4_glm_router_better(scores, a, b);
|
||||
if (swap) { idx[tid] = b; idx[other] = a; }
|
||||
}
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
if (lane < 8u) token_selected[lane] = idx[lane];
|
||||
} else {
|
||||
uint available = 511u;
|
||||
for (uint rank = 0u; rank < 8u; ++rank) {
|
||||
float best_score = -INFINITY;
|
||||
uint best_id = 0xffffffffu;
|
||||
for (uint i = 0u; i < 9u; ++i) {
|
||||
const uint expert = lane + 32u * i;
|
||||
if ((available & (1u << i)) &&
|
||||
(values[i] > best_score || (values[i] == best_score && expert < best_id))) {
|
||||
best_score = values[i];
|
||||
best_id = expert;
|
||||
}
|
||||
}
|
||||
const float group_score = simd_max(best_score);
|
||||
const uint group_id = simd_min(best_score == group_score ? best_id : 0xffffffffu);
|
||||
if (lane == rank) token_selected[rank] = (int32_t)group_id;
|
||||
if ((group_id & 31u) == lane) available &= ~(1u << (group_id >> 5u));
|
||||
}
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_device);
|
||||
const uint k_used = min(args.n_expert_used, args.n_expert);
|
||||
if (lane < k_used) {
|
||||
float sum = 0.0f;
|
||||
for (uint i = 0; i < k_used; i++) sum += token_probs[(uint)token_selected[i]];
|
||||
sum = max(sum, 6.103515625e-5f);
|
||||
token_weights[lane] = token_probs[(uint)token_selected[lane]] / sum * args.expert_weight_scale;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_glm_router_select_top8(
|
||||
constant ds4_metal_args_glm_router_select_one & args,
|
||||
device const float *logits,
|
||||
device const float *bias,
|
||||
device int32_t *selected,
|
||||
device float *weights,
|
||||
device float *probs,
|
||||
threadgroup float *scratch [[threadgroup(0)]],
|
||||
uint token [[threadgroup_position_in_grid]],
|
||||
uint lane [[thread_index_in_simdgroup]]) {
|
||||
ds4_glm_router_select_top8_impl(args, logits, bias, selected, weights,
|
||||
probs, scratch, token, lane);
|
||||
}
|
||||
|
||||
|
||||
// Adapted from IngeniousIdiocy/ds4 95eb218, MIT: one grid carries the
|
||||
// F32 GLM router and two virtual four-SIMDgroup shared-expert cohorts.
|
||||
// The current top-eight body is retained, including its exceptional fallback.
|
||||
kernel void kernel_glm53_router_shared_exact(
|
||||
constant ds4_metal_args_mul_mv & rargs,
|
||||
constant ds4_metal_args_mul_mv & sargs,
|
||||
constant ds4_metal_args_glm_router_select_one & args,
|
||||
device const char *rw, device const char *gw, device const char *uw,
|
||||
device const char *x, device char *logits, device const float *bias,
|
||||
device int32_t *selected, device float *weights, device float *probs,
|
||||
device atomic_uint *counter, device char *mid,
|
||||
constant float &clamp_value,
|
||||
threadgroup char *scratch [[threadgroup(0)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
if (tg.x >= 144u) {
|
||||
const ushort cohort = sg >> 2;
|
||||
uint3 vtg = tg;
|
||||
vtg.x = (tg.x - 144u) * 2u + cohort;
|
||||
kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<2, false, 4>(
|
||||
sargs, gw, uw, x, mid, mid, mid, clamp_value,
|
||||
scratch + cohort * 512u, vtg, lane, sg & 3u);
|
||||
return;
|
||||
}
|
||||
kernel_mul_mv_t_t_4_impl<float, float4, float, float4, 2,
|
||||
constant ds4_metal_args_mul_mv &, true>(
|
||||
rargs, rw, x, logits, scratch, tg, lane, sg);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device);
|
||||
atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
|
||||
threadgroup uint elected;
|
||||
if (tid == 0u) {
|
||||
elected = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == 143u;
|
||||
if (elected) atomic_store_explicit(counter, 0u, memory_order_relaxed);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (!elected) return;
|
||||
atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
|
||||
device atomic_uint *slots = (device atomic_uint *)logits;
|
||||
for (uint i = tid; i < 288u; i += 256u)
|
||||
((device float *)logits)[i] = as_type<float>(atomic_load_explicit(slots+i, memory_order_relaxed));
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
if (sg == 0u)
|
||||
ds4_glm_router_select_top8_impl(args, (device const float *)logits,
|
||||
bias, selected, weights, probs, (threadgroup float *)scratch, 0u, lane);
|
||||
}
|
||||
|
||||
// Batched Flash-router weight finalization after selection is already known.
|
||||
// Six active lanes deliberately match kernel_sum_rows_f32_f32's reduction
|
||||
// topology. The denominator and divided weights cross threadgroup storage
|
||||
|
||||
+217
-16
@@ -26,24 +26,20 @@ kernel void kernel_glm53_embedding_bf16(
|
||||
: 0.0f;
|
||||
}
|
||||
|
||||
static inline void glm53_mul_mv_bf16_f32_row(
|
||||
constant glm53_bf16_matmul_args &args,
|
||||
/* The accumulation, split out unchanged so an epilogue kernel can use the sum
|
||||
* before it is stored. Callers must range-check out_row and token first. */
|
||||
static inline float glm53_mul_mv_bf16_f32_row_sum(
|
||||
uint in_dim,
|
||||
device const ushort *weights,
|
||||
device const float *x,
|
||||
device float *out,
|
||||
uint2 tgpig,
|
||||
ushort lane,
|
||||
ushort sg,
|
||||
ushort nsg) {
|
||||
const uint out_row = tgpig.x * (uint)nsg + sg;
|
||||
const uint token = tgpig.y;
|
||||
if (out_row >= args.out_dim || token >= args.n_rows) return;
|
||||
|
||||
device const ushort *w = weights + (ulong)out_row * args.in_dim;
|
||||
device const float *xr = x + (ulong)token * args.in_dim;
|
||||
uint out_row,
|
||||
uint token,
|
||||
ushort lane) {
|
||||
device const ushort *w = weights + (ulong)out_row * in_dim;
|
||||
device const float *xr = x + (ulong)token * in_dim;
|
||||
float sum = 0.0f;
|
||||
uint k = lane;
|
||||
for (; k + 224u < args.in_dim; k += 256u) {
|
||||
for (; k + 224u < in_dim; k += 256u) {
|
||||
const ushort w0 = w[k];
|
||||
const ushort w1 = w[k + 32u];
|
||||
const ushort w2 = w[k + 64u];
|
||||
@@ -69,10 +65,26 @@ static inline void glm53_mul_mv_bf16_f32_row(
|
||||
sum = fma(glm53_bf16_to_f32(w6), x6, sum);
|
||||
sum = fma(glm53_bf16_to_f32(w7), x7, sum);
|
||||
}
|
||||
for (; k < args.in_dim; k += 32u) {
|
||||
for (; k < in_dim; k += 32u) {
|
||||
sum = fma(glm53_bf16_to_f32(w[k]), xr[k], sum);
|
||||
}
|
||||
sum = simd_sum(sum);
|
||||
return simd_sum(sum);
|
||||
}
|
||||
|
||||
static inline void glm53_mul_mv_bf16_f32_row(
|
||||
constant glm53_bf16_matmul_args &args,
|
||||
device const ushort *weights,
|
||||
device const float *x,
|
||||
device float *out,
|
||||
uint2 tgpig,
|
||||
ushort lane,
|
||||
ushort sg,
|
||||
ushort nsg) {
|
||||
const uint out_row = tgpig.x * (uint)nsg + sg;
|
||||
const uint token = tgpig.y;
|
||||
if (out_row >= args.out_dim || token >= args.n_rows) return;
|
||||
const float sum =
|
||||
glm53_mul_mv_bf16_f32_row_sum(args.in_dim, weights, x, out_row, token, lane);
|
||||
if (lane == 0u) out[(ulong)token * args.out_dim + out_row] = sum;
|
||||
}
|
||||
|
||||
@@ -91,6 +103,136 @@ kernel void kernel_glm53_mul_mv_bf16_f32(
|
||||
tgpig, lane, sg, nsg);
|
||||
}
|
||||
|
||||
/*
|
||||
* BF16 matvec with the HC expansion folded into its epilogue.
|
||||
*
|
||||
* The simdgroup that finishes output row d already holds that row's value in
|
||||
* lane 0, so it can expand it into the four HC streams there instead of
|
||||
* writing it out and having a second dispatch read it straight back. This is
|
||||
* the shape kernel_dsv4_q8_hc_expand4_q8_0 already uses for DeepSeek, in BF16.
|
||||
*
|
||||
* Decode only: one token, HC = 4. The arithmetic and the operand order match
|
||||
* kernel_dsv4_hc_expand4 exactly, including that comb is indexed [j][h].
|
||||
*/
|
||||
kernel void kernel_glm53_mul_mv_bf16_f32_hc_expand4(
|
||||
constant glm53_bf16_matmul_args &args,
|
||||
device const ushort *weights,
|
||||
device const float *x,
|
||||
device float *out,
|
||||
device const float *residual,
|
||||
device const float *post,
|
||||
device const float *comb,
|
||||
device float *hc_out,
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort nsg [[simdgroups_per_threadgroup]]) {
|
||||
const uint out_row = tgpig.x * (uint)nsg + sg;
|
||||
const uint token = tgpig.y;
|
||||
if (out_row >= args.out_dim || token >= args.n_rows) return;
|
||||
const float sum =
|
||||
glm53_mul_mv_bf16_f32_row_sum(args.in_dim, weights, x, out_row, token, lane);
|
||||
if (lane != 0u) return;
|
||||
out[(ulong)token * args.out_dim + out_row] = sum;
|
||||
|
||||
const uint n = args.out_dim;
|
||||
const float r0 = residual[0u * n + out_row];
|
||||
const float r1 = residual[1u * n + out_row];
|
||||
const float r2 = residual[2u * n + out_row];
|
||||
const float r3 = residual[3u * n + out_row];
|
||||
for (uint h = 0u; h < 4u; ++h) {
|
||||
float acc = sum * post[h];
|
||||
acc += comb[0u * 4u + h] * r0;
|
||||
acc += comb[1u * 4u + h] * r1;
|
||||
acc += comb[2u * 4u + h] * r2;
|
||||
acc += comb[3u * 4u + h] * r3;
|
||||
hc_out[h * n + out_row] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
struct glm53_bf16_trio_args {
|
||||
uint in_dim;
|
||||
uint out_dim_ab;
|
||||
uint out_dim_c;
|
||||
uint n_rows;
|
||||
};
|
||||
|
||||
/*
|
||||
* Three matvecs over one shared input in a single dispatch, where the third
|
||||
* has a shorter output than the first two. GLM 5.3's KDA gate chain is
|
||||
* exactly that shape: f_a and g_a are [4096 -> 128] and beta is [4096 -> 64],
|
||||
* all reading attn_norm. The pair kernel could not carry beta because it
|
||||
* assumes one output width for every slot.
|
||||
*
|
||||
* The grid is sized for the wider pair, so the beta slot's upper threadgroups
|
||||
* exit on the bounds check.
|
||||
*/
|
||||
kernel void kernel_glm53_mul_mv_bf16_f32_trio(
|
||||
constant glm53_bf16_trio_args &args,
|
||||
device const ushort *weights_a,
|
||||
device const ushort *weights_b,
|
||||
device const ushort *weights_c,
|
||||
device const float *x,
|
||||
device float *out_a,
|
||||
device float *out_b,
|
||||
device float *out_c,
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort nsg [[simdgroups_per_threadgroup]]) {
|
||||
const uint slot = tgpig.z;
|
||||
device const ushort *w = slot == 0u ? weights_a
|
||||
: (slot == 1u ? weights_b : weights_c);
|
||||
device float *out = slot == 0u ? out_a : (slot == 1u ? out_b : out_c);
|
||||
const uint out_dim = slot == 2u ? args.out_dim_c : args.out_dim_ab;
|
||||
const uint out_row = tgpig.x * (uint)nsg + sg;
|
||||
const uint token = tgpig.y;
|
||||
if (out_row >= out_dim || token >= args.n_rows) return;
|
||||
const float sum =
|
||||
glm53_mul_mv_bf16_f32_row_sum(args.in_dim, w, x, out_row, token, lane);
|
||||
if (lane == 0u) out[(ulong)token * out_dim + out_row] = sum;
|
||||
}
|
||||
|
||||
/* Flatten the six independent KDA projections into one grid, avoiding empty
|
||||
* groups for the three short outputs. Each row retains the existing dot. */
|
||||
kernel void kernel_glm53_kda_inputs_bf16(
|
||||
device const ushort *weights_q,
|
||||
device const ushort *weights_k,
|
||||
device const ushort *weights_v,
|
||||
device const ushort *weights_f,
|
||||
device const ushort *weights_g,
|
||||
device const ushort *weights_beta,
|
||||
device const float *x,
|
||||
device float *out_q,
|
||||
device float *out_k,
|
||||
device float *out_v,
|
||||
device float *out_f,
|
||||
device float *out_g,
|
||||
device float *out_beta,
|
||||
uint group [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
uint row = group * 4u + sg;
|
||||
if (row >= 24896u) return;
|
||||
device const ushort *weights;
|
||||
device float *out;
|
||||
if (row < 8192u) {
|
||||
weights = weights_q; out = out_q;
|
||||
} else if (row < 16384u) {
|
||||
row -= 8192u; weights = weights_k; out = out_k;
|
||||
} else if (row < 24576u) {
|
||||
row -= 16384u; weights = weights_v; out = out_v;
|
||||
} else if (row < 24704u) {
|
||||
row -= 24576u; weights = weights_f; out = out_f;
|
||||
} else if (row < 24832u) {
|
||||
row -= 24704u; weights = weights_g; out = out_g;
|
||||
} else {
|
||||
row -= 24832u; weights = weights_beta; out = out_beta;
|
||||
}
|
||||
const float sum = glm53_mul_mv_bf16_f32_row_sum(4096u, weights, x, row, 0u, lane);
|
||||
if (lane == 0u) out[row] = sum;
|
||||
}
|
||||
|
||||
kernel void kernel_glm53_mul_mv_bf16_f32_qkv(
|
||||
constant glm53_bf16_matmul_args &args,
|
||||
device const ushort *weights_q,
|
||||
@@ -112,6 +254,32 @@ kernel void kernel_glm53_mul_mv_bf16_f32_qkv(
|
||||
tgpig.xy, lane, sg, nsg);
|
||||
}
|
||||
|
||||
/*
|
||||
* Two independent matvecs of the same shape in one dispatch, selected by
|
||||
* tgpig.z, exactly as the qkv variant above selects three. The inputs are
|
||||
* separate pointers rather than one shared row, which lets this serve both
|
||||
* halves of the GLM 5.3 KDA gate chain: f_a/g_a read the same attn_norm row,
|
||||
* while f_b/g_b read the two different low-rank vectors those produce.
|
||||
*/
|
||||
kernel void kernel_glm53_mul_mv_bf16_f32_pair(
|
||||
constant glm53_bf16_matmul_args &args,
|
||||
device const ushort *weights_a,
|
||||
device const ushort *weights_b,
|
||||
device const float *x_a,
|
||||
device const float *x_b,
|
||||
device float *out_a,
|
||||
device float *out_b,
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort nsg [[simdgroups_per_threadgroup]]) {
|
||||
device const ushort *weights = tgpig.z == 0u ? weights_a : weights_b;
|
||||
device const float *x = tgpig.z == 0u ? x_a : x_b;
|
||||
device float *out = tgpig.z == 0u ? out_a : out_b;
|
||||
glm53_mul_mv_bf16_f32_row(args, weights, x, out,
|
||||
tgpig.xy, lane, sg, nsg);
|
||||
}
|
||||
|
||||
struct glm53_bf16_block16 {
|
||||
ushort v[16];
|
||||
};
|
||||
@@ -141,3 +309,36 @@ kernel glm53_mul_mm_bf16_t kernel_mul_mm<
|
||||
half, half2x4, simdgroup_half8x8,
|
||||
glm53_bf16_block16, 1, glm53_dequantize_bf16,
|
||||
half, half4x4, float, float2x4>;
|
||||
|
||||
/* Mixed Q8 QKV + BF16 gate inputs share a dispatch, retaining each original
|
||||
* projection helper and its reduction order. Inspired by the M3 Ultra fork's
|
||||
* projection fold; this accepts our existing mixed model without requantizing. */
|
||||
kernel void kernel_glm53_kda_inputs_q8_bf16(
|
||||
device const char *wq, device const char *wk, device const char *wv,
|
||||
device const ushort *wf, device const ushort *wg, device const ushort *wb,
|
||||
device const float *x,
|
||||
device float *oq, device float *ok, device float *ov,
|
||||
device float *of, device float *og, device float *ob,
|
||||
constant ds4_metal_args_mul_mv &args,
|
||||
threadgroup char *scratch [[threadgroup(0)]],
|
||||
uint group [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
if (group < 12288u) {
|
||||
const uint slot=group/4096u;
|
||||
device const char *w=slot==0u?wq:slot==1u?wk:wv;
|
||||
device float *o=slot==0u?oq:slot==1u?ok:ov;
|
||||
kernel_mul_mv_q8_0_f32_impl<2, constant ds4_metal_args_mul_mv &>(
|
||||
args,w,(device const char *)x,(device char *)o,scratch,
|
||||
uint3(group%4096u,0u,0u),lane,sg);
|
||||
} else {
|
||||
uint row=(group-12288u)*4u+sg;
|
||||
if (row>=320u) return;
|
||||
device const ushort *w; device float *o;
|
||||
if(row<128u){w=wf;o=of;}
|
||||
else if(row<256u){row-=128u;w=wg;o=og;}
|
||||
else{row-=256u;w=wb;o=ob;}
|
||||
const float sum=glm53_mul_mv_bf16_f32_row_sum(4096u,w,x,row,0u,lane);
|
||||
if(lane==0u)o[row]=sum;
|
||||
}
|
||||
}
|
||||
+294
-21
@@ -8,10 +8,11 @@ struct glm53_kda_args {
|
||||
};
|
||||
|
||||
/*
|
||||
* One threadgroup owns one (sequence, head). Four simdgroups update four
|
||||
* value rows concurrently; every lane owns four adjacent key columns.
|
||||
* One threadgroup owns one (sequence, head). Four SIMDgroups own disjoint
|
||||
* value rows; every lane owns four adjacent key columns.
|
||||
*/
|
||||
kernel void kernel_glm53_kda_decode(
|
||||
template<uint VALUES>
|
||||
kernel void kernel_glm53_kda_decode_values(
|
||||
constant glm53_kda_args &args,
|
||||
device const float *q_in,
|
||||
device const float *k_in,
|
||||
@@ -48,6 +49,7 @@ kernel void kernel_glm53_kda_decode(
|
||||
threadgroup float *reduce_k = reduce_q + 4u;
|
||||
threadgroup float *reduce_o = reduce_k + 4u;
|
||||
threadgroup float *beta_shared = reduce_o + 4u;
|
||||
threadgroup float *a_decay_shared = beta_shared + 1u;
|
||||
|
||||
const uint projection = args.n_heads * D;
|
||||
const uint channel = head * D + tid;
|
||||
@@ -90,16 +92,18 @@ kernel void kernel_glm53_kda_decode(
|
||||
sq[tid] = q_acc / (1.0f + exp(-q_acc));
|
||||
sk[tid] = k_acc / (1.0f + exp(-k_acc));
|
||||
sv[tid] = v_acc / (1.0f + exp(-v_acc));
|
||||
const float gate = raw_gate[input_base + tid] + dt_bias[channel];
|
||||
sd[tid] = exp(args.lower_bound *
|
||||
(1.0f / (1.0f + exp(-exp(a_log[head]) * gate))));
|
||||
}
|
||||
if (tid == 0u) {
|
||||
beta_shared[0] =
|
||||
1.0f / (1.0f + exp(-raw_beta[(ulong)row * args.n_heads + head]));
|
||||
/* head is uniform over the threadgroup, so exp(a_log[head]) is a
|
||||
* single value; every one of the D channels used to recompute it. */
|
||||
a_decay_shared[0] = exp(a_log[head]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup |
|
||||
mem_flags::mem_device);
|
||||
/* Only threadgroup memory is shared between threads here: the conv-state
|
||||
* writes above are each thread's own channel and no thread reads another's,
|
||||
* so the barrier does not need device scope. */
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float q_sumsq = sq[tid] * sq[tid];
|
||||
float k_sumsq = sk[tid] * sk[tid];
|
||||
@@ -119,6 +123,9 @@ kernel void kernel_glm53_kda_decode(
|
||||
if (tid < D) {
|
||||
sq[tid] *= q_scale;
|
||||
sk[tid] *= k_scale;
|
||||
const float gate = raw_gate[input_base + tid] + dt_bias[channel];
|
||||
sd[tid] = exp(args.lower_bound *
|
||||
(1.0f / (1.0f + exp(-a_decay_shared[0] * gate))));
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
@@ -129,20 +136,49 @@ kernel void kernel_glm53_kda_decode(
|
||||
const ulong state_head =
|
||||
((ulong)row * args.n_heads + head) * D * D;
|
||||
|
||||
for (uint value = sg; value < D; value += 4u) {
|
||||
device float4 *hptr =
|
||||
(device float4 *)(state + state_head + (ulong)value * D + k0);
|
||||
float4 h = *hptr * decay4;
|
||||
float hk = dot(h, k4);
|
||||
hk = simd_sum(hk);
|
||||
const float delta_v = (sv[value] - hk) * beta_shared[0];
|
||||
h = fma(k4, float4(delta_v), h);
|
||||
*hptr = h;
|
||||
float hq = simd_sum(dot(h, q4));
|
||||
if (lane == 0u) so[value] = hq;
|
||||
/* Interleave independent value rows to hide load/reduction latency.
|
||||
* Every row retains its original SIMDgroup, dot/FMA sequence, and stores;
|
||||
* normalization still uses the same four-SIMDgroup reduction tree. */
|
||||
if (VALUES == 1u) {
|
||||
for (uint value = sg; value < D; value += 4u) {
|
||||
device float4 *hptr =
|
||||
(device float4 *)(state + state_head + (ulong)value * D + k0);
|
||||
float4 h = *hptr * decay4;
|
||||
float hk = dot(h, k4);
|
||||
hk = simd_sum(hk);
|
||||
const float delta_v = (sv[value] - hk) * beta_shared[0];
|
||||
h = fma(k4, float4(delta_v), h);
|
||||
*hptr = h;
|
||||
float hq = simd_sum(dot(h, q4));
|
||||
if (lane == 0u) so[value] = hq;
|
||||
}
|
||||
} else {
|
||||
for (uint value = sg; value < D; value += 4u * VALUES) {
|
||||
device float4 *hptr[VALUES];
|
||||
float4 h[VALUES];
|
||||
float hk[VALUES];
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
hptr[i] = (device float4 *)(state + state_head + (ulong)(value + i * 4u) * D + k0);
|
||||
h[i] = *hptr[i] * decay4;
|
||||
}
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
hk[i] = dot(h[i], k4);
|
||||
hk[i] = simd_sum(hk[i]);
|
||||
}
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
const float delta_v = (sv[value + i * 4u] - hk[i]) * beta_shared[0];
|
||||
h[i] = fma(k4, float4(delta_v), h[i]);
|
||||
*hptr[i] = h[i];
|
||||
}
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
const float hq = simd_sum(dot(h[i], q4));
|
||||
if (lane == 0u) so[value + i * 4u] = hq;
|
||||
}
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup |
|
||||
mem_flags::mem_device);
|
||||
/* Likewise: the state writes above are not re-read in this kernel, only
|
||||
* so[] crosses simdgroups. */
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float o_sumsq = so[tid] * so[tid];
|
||||
o_sumsq = simd_sum(o_sumsq);
|
||||
@@ -159,6 +195,12 @@ kernel void kernel_glm53_kda_decode(
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_glm53_kda_decode_values<1u>) glm53_kda_decode_values_t;
|
||||
template [[host_name("kernel_glm53_kda_decode")]]
|
||||
kernel glm53_kda_decode_values_t kernel_glm53_kda_decode_values<1u>;
|
||||
template [[host_name("kernel_glm53_kda_decode_v4")]]
|
||||
kernel glm53_kda_decode_values_t kernel_glm53_kda_decode_values<4u>;
|
||||
|
||||
kernel void kernel_glm53_kda_prefill_prepare(
|
||||
constant glm53_kda_args &args,
|
||||
device float *q,
|
||||
@@ -246,6 +288,184 @@ kernel void kernel_glm53_kda_prefill_prepare(
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Boundary rows for the blocked prepare kernel below.
|
||||
*
|
||||
* A block starts its causal convolution window on the raw q/k/v of the three
|
||||
* rows before it, and the block before it overwrites exactly those rows with
|
||||
* its normalized outputs. Block 0 also needs an immutable copy of the incoming
|
||||
* conv state: the last block may overwrite that state before block 0 starts.
|
||||
*/
|
||||
struct glm53_kda_blocked_args {
|
||||
uint n_heads;
|
||||
uint n_rows;
|
||||
uint block_rows;
|
||||
uint n_blocks;
|
||||
uint block_base;
|
||||
float lower_bound;
|
||||
float norm_eps;
|
||||
};
|
||||
|
||||
kernel void kernel_glm53_kda_prefill_conv_halo(
|
||||
constant glm53_kda_blocked_args &args,
|
||||
device const float *q,
|
||||
device const float *k,
|
||||
device const float *v,
|
||||
device float *halo,
|
||||
device const float *conv_state,
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]]) {
|
||||
constexpr uint D = 128u;
|
||||
constexpr uint HISTORY = 3u;
|
||||
constexpr uint NTH = 256u; /* matches the dispatch */
|
||||
const uint projection = args.n_heads * D;
|
||||
const uint block = tgpig.x;
|
||||
const uint w = tgpig.y;
|
||||
if (block >= args.n_blocks || w >= HISTORY) return;
|
||||
const ulong plane = (ulong)args.n_blocks * HISTORY * projection;
|
||||
const ulong slot = ((ulong)block * HISTORY + w) * projection;
|
||||
for (uint c = tid; c < projection; c += NTH) {
|
||||
if (block == 0u) {
|
||||
const ulong index = (ulong)w * projection + c;
|
||||
halo[slot + c] = conv_state[index];
|
||||
halo[plane + slot + c] = conv_state[HISTORY * projection + index];
|
||||
halo[2u * plane + slot + c] = conv_state[2u * HISTORY * projection + index];
|
||||
} else {
|
||||
const uint token = block * args.block_rows + w - HISTORY;
|
||||
const ulong index = (ulong)token * projection + c;
|
||||
halo[slot + c] = q[index];
|
||||
halo[plane + slot + c] = k[index];
|
||||
halo[2u * plane + slot + c] = v[index];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Token-parallel form of kernel_glm53_kda_prefill_prepare.
|
||||
*
|
||||
* The serial kernel runs one threadgroup per head -- 64 of them on an 80-core
|
||||
* GPU -- and walks all the chunk's tokens inside it, so it costs 3.4 ms of a
|
||||
* 2048-token layer at about 5% occupancy. Nothing in it is actually
|
||||
* sequential: the causal convolution reads the raw q/k/v of t-3..t, which are
|
||||
* all known before the kernel starts.
|
||||
*
|
||||
* One threadgroup now owns (block of block_rows tokens, head) and keeps the
|
||||
* three-row convolution history in registers instead of re-reading and
|
||||
* re-writing the device conv state every token. Its first three rows come
|
||||
* from the immutable halo above, including the incoming state for block 0.
|
||||
* The last block leaves the outgoing conv state exactly where the serial kernel
|
||||
* left it.
|
||||
*
|
||||
* Every value keeps the serial kernel's expression and order: the same
|
||||
* four-term fma chain in the same order, the same silu, the same
|
||||
* 4-simdgroup RMS reduction, the same decay-gate expression.
|
||||
*/
|
||||
kernel void kernel_glm53_kda_prefill_prepare_blocked(
|
||||
constant glm53_kda_blocked_args &args,
|
||||
device float *q,
|
||||
device float *k,
|
||||
device float *v,
|
||||
device float *raw_gate,
|
||||
device const float *q_conv,
|
||||
device const float *k_conv,
|
||||
device const float *v_conv,
|
||||
device const float *a_log,
|
||||
device const float *dt_bias,
|
||||
device float *conv_state,
|
||||
device const float *halo,
|
||||
threadgroup float *scratch [[threadgroup(0)]],
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr uint D = 128u;
|
||||
constexpr uint HISTORY = 3u;
|
||||
const uint block = tgpig.x + args.block_base;
|
||||
const uint head = tgpig.y;
|
||||
if (block >= args.n_blocks || head >= args.n_heads) return;
|
||||
threadgroup float *sq = scratch;
|
||||
threadgroup float *sk = sq + D;
|
||||
threadgroup float *reduce_q = sk + D;
|
||||
threadgroup float *reduce_k = reduce_q + 4u;
|
||||
const uint projection = args.n_heads * D;
|
||||
const uint channel = head * D + tid;
|
||||
const uint token0 = block * args.block_rows;
|
||||
const uint token_end = token0 + min(args.block_rows, args.n_rows - token0);
|
||||
if (token0 >= token_end) return;
|
||||
|
||||
float hq[HISTORY];
|
||||
float hk[HISTORY];
|
||||
float hv[HISTORY];
|
||||
const ulong plane = (ulong)args.n_blocks * HISTORY * projection;
|
||||
const ulong slot = (ulong)block * HISTORY * projection + channel;
|
||||
for (uint w = 0; w < HISTORY; w++) {
|
||||
hq[w] = halo[slot + w * projection];
|
||||
hk[w] = halo[plane + slot + w * projection];
|
||||
hv[w] = halo[2u * plane + slot + w * projection];
|
||||
}
|
||||
|
||||
for (uint token = token0; token < token_end; token++) {
|
||||
const ulong index = (ulong)token * projection + channel;
|
||||
float q_acc = 0.0f;
|
||||
float k_acc = 0.0f;
|
||||
float v_acc = 0.0f;
|
||||
for (uint w = 0; w < HISTORY; w++) {
|
||||
q_acc = fma(hq[w], q_conv[(ulong)channel * 4u + w], q_acc);
|
||||
k_acc = fma(hk[w], k_conv[(ulong)channel * 4u + w], k_acc);
|
||||
v_acc = fma(hv[w], v_conv[(ulong)channel * 4u + w], v_acc);
|
||||
}
|
||||
const float q_new = q[index];
|
||||
const float k_new = k[index];
|
||||
const float v_new = v[index];
|
||||
q_acc = fma(q_new, q_conv[(ulong)channel * 4u + 3u], q_acc);
|
||||
k_acc = fma(k_new, k_conv[(ulong)channel * 4u + 3u], k_acc);
|
||||
v_acc = fma(v_new, v_conv[(ulong)channel * 4u + 3u], v_acc);
|
||||
hq[0] = hq[1]; hq[1] = hq[2]; hq[2] = q_new;
|
||||
hk[0] = hk[1]; hk[1] = hk[2]; hk[2] = k_new;
|
||||
hv[0] = hv[1]; hv[1] = hv[2]; hv[2] = v_new;
|
||||
|
||||
sq[tid] = q_acc / (1.0f + exp(-q_acc));
|
||||
sk[tid] = k_acc / (1.0f + exp(-k_acc));
|
||||
v[index] = v_acc / (1.0f + exp(-v_acc));
|
||||
const float gate = raw_gate[index] + dt_bias[channel];
|
||||
raw_gate[index] = exp(args.lower_bound *
|
||||
(1.0f / (1.0f + exp(-exp(a_log[head]) * gate))));
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup |
|
||||
mem_flags::mem_device);
|
||||
|
||||
float q_sumsq = simd_sum(sq[tid] * sq[tid]);
|
||||
float k_sumsq = simd_sum(sk[tid] * sk[tid]);
|
||||
if (lane == 0u) {
|
||||
reduce_q[sg] = q_sumsq;
|
||||
reduce_k[sg] = k_sumsq;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
float q_total = lane < 4u ? reduce_q[lane] : 0.0f;
|
||||
float k_total = lane < 4u ? reduce_k[lane] : 0.0f;
|
||||
q_total = simd_sum(q_total);
|
||||
k_total = simd_sum(k_total);
|
||||
q[index] = sq[tid] * rsqrt(q_total + 1.0e-6f) *
|
||||
0x1.6a09e6p-4f;
|
||||
k[index] = sk[tid] * rsqrt(k_total + 1.0e-6f);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup |
|
||||
mem_flags::mem_device);
|
||||
}
|
||||
|
||||
/* The serial kernel leaves the conv state holding the raw q/k/v of the
|
||||
* last three rows it processed, which is what this block's history is
|
||||
* once its last token has shifted through. */
|
||||
if (token_end == args.n_rows) {
|
||||
device float *q_state = conv_state;
|
||||
device float *k_state = q_state + HISTORY * projection;
|
||||
device float *v_state = k_state + HISTORY * projection;
|
||||
for (uint w = 0; w < HISTORY; w++) {
|
||||
q_state[(ulong)w * projection + channel] = hq[w];
|
||||
k_state[(ulong)w * projection + channel] = hk[w];
|
||||
v_state[(ulong)w * projection + channel] = hv[w];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_glm53_kda_prefill_recurrence(
|
||||
constant glm53_kda_args &args,
|
||||
device const float *q,
|
||||
@@ -286,6 +506,59 @@ kernel void kernel_glm53_kda_prefill_recurrence(
|
||||
*state_ptr = h;
|
||||
}
|
||||
|
||||
/* Each value row is an independent recurrence. Carrying two or four rows in
|
||||
* a SIMDgroup reuses q/k/decay and beta across them, while keeping the serial
|
||||
* token order, dot reductions and FMA expression of the reference above. */
|
||||
template <uint VALUES>
|
||||
kernel void kernel_glm53_kda_prefill_recurrence_values(
|
||||
constant glm53_kda_args &args,
|
||||
device const float *q,
|
||||
device const float *k,
|
||||
device const float *v,
|
||||
device const float *decay,
|
||||
device const float *raw_beta,
|
||||
device float *state,
|
||||
device float *out,
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr uint D = 128u;
|
||||
const uint head = tgpig.x;
|
||||
const uint value0 = (tgpig.y * 4u + sg) * VALUES;
|
||||
if (head >= args.n_heads || value0 + VALUES > D) return;
|
||||
const uint projection = args.n_heads * D;
|
||||
const uint k0 = lane * 4u;
|
||||
float4 h[VALUES];
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
h[i] = *((device float4 *)(state + ((ulong)head * D + value0 + i) * D + k0));
|
||||
}
|
||||
for (uint token = 0; token < args.n_rows; token++) {
|
||||
const ulong base = (ulong)token * projection + head * D;
|
||||
const float4 q4 = *((device const float4 *)(q + base + k0));
|
||||
const float4 k4 = *((device const float4 *)(k + base + k0));
|
||||
const float4 decay4 = *((device const float4 *)(decay + base + k0));
|
||||
const float beta = 1.0f /
|
||||
(1.0f + exp(-raw_beta[(ulong)token * args.n_heads + head]));
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
h[i] *= decay4;
|
||||
const float hk = simd_sum(dot(h[i], k4));
|
||||
const float delta_v = (v[base + value0 + i] - hk) * beta;
|
||||
h[i] = fma(k4, float4(delta_v), h[i]);
|
||||
const float result = simd_sum(dot(h[i], q4));
|
||||
if (lane == 0u) out[base + value0 + i] = result;
|
||||
}
|
||||
}
|
||||
FOR_UNROLL (uint i = 0; i < VALUES; i++) {
|
||||
*((device float4 *)(state + ((ulong)head * D + value0 + i) * D + k0)) = h[i];
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_glm53_kda_prefill_recurrence_values<2u>) glm53_kda_recurrence_values_t;
|
||||
template [[host_name("kernel_glm53_kda_prefill_recurrence_v2")]]
|
||||
kernel glm53_kda_recurrence_values_t kernel_glm53_kda_prefill_recurrence_values<2u>;
|
||||
template [[host_name("kernel_glm53_kda_prefill_recurrence_v4")]]
|
||||
kernel glm53_kda_recurrence_values_t kernel_glm53_kda_prefill_recurrence_values<4u>;
|
||||
|
||||
kernel void kernel_glm53_kda_prefill_output(
|
||||
constant glm53_kda_args &args,
|
||||
device float *out,
|
||||
|
||||
@@ -8825,6 +8825,8 @@ template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull")
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id;
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs;
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_tail_cull;
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4, true>) mul_mm_id_q4_K_f16_rhs_tail_cull;
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4, true>) mul_mm_id_q4_K_f32_rhs_tail_cull;
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4>) mul_mm_id_mxfp4_f16_rhs_half_lut;
|
||||
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_half_lut_tail_cull;
|
||||
typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_addr;
|
||||
@@ -8834,6 +8836,7 @@ typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, hal
|
||||
template [[host_name("kernel_mul_mm_id_q8_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q2_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_K_f32_tail_cull")]] kernel mul_mm_id_q4_K_f32_rhs_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4, true>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q6_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>;
|
||||
@@ -8841,6 +8844,7 @@ template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id ker
|
||||
template [[host_name("kernel_mul_mm_id_q8_0_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_K_f16_tail_cull")]] kernel mul_mm_id_q4_K_f16_rhs_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4, true>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q6_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, half, half4x4, half, half2x4>;
|
||||
|
||||
+25
-7
@@ -5906,7 +5906,8 @@ static bool test_logprob_vector_case_disabled(const char *path,
|
||||
|
||||
static void test_official_logprob_vectors_run(const char *case_filter) {
|
||||
const char *path = getenv("DS4_TEST_VECTOR_FILE");
|
||||
if (!path || !path[0]) {
|
||||
const bool default_fixture = !path || !path[0];
|
||||
if (default_fixture) {
|
||||
path = "tests/test-vectors/flash-0731/official.vec";
|
||||
}
|
||||
FILE *fp = fopen(path, "rb");
|
||||
@@ -5934,7 +5935,15 @@ static void test_official_logprob_vectors_run(const char *case_filter) {
|
||||
|
||||
test_vec_case vc;
|
||||
int ran = 0;
|
||||
while (test_read_vector_case(fp, &vc)) {
|
||||
/* The default vectors describe DeepSeek V4 Flash, not GLM. An explicitly
|
||||
* selected fixture remains available for model-specific comparisons. */
|
||||
const bool compatible_fixture = !default_fixture || !ds4_engine_is_glm_dsa(engine);
|
||||
if (!compatible_fixture) {
|
||||
fprintf(stderr, "ds4-test: DeepSeek API vectors skipped for %s; "
|
||||
"set DS4_TEST_VECTOR_FILE to a matching fixture\n",
|
||||
ds4_engine_model_name(engine));
|
||||
}
|
||||
while (compatible_fixture && test_read_vector_case(fp, &vc)) {
|
||||
if (!test_fill_vector_case(fp, &vc)) break;
|
||||
if (case_filter && case_filter[0] && strcmp(vc.id, case_filter)) {
|
||||
continue;
|
||||
@@ -5948,7 +5957,7 @@ static void test_official_logprob_vectors_run(const char *case_filter) {
|
||||
test_logprob_vector_case(engine, &vc);
|
||||
ran++;
|
||||
}
|
||||
TEST_ASSERT(!case_filter || !case_filter[0] || ran == 1);
|
||||
TEST_ASSERT(!compatible_fixture || !case_filter || !case_filter[0] || ran == 1);
|
||||
ds4_engine_close(engine);
|
||||
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
|
||||
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
|
||||
@@ -6090,12 +6099,13 @@ static int test_local_golden_overlap(const test_local_golden_case *tc,
|
||||
|
||||
static float test_local_golden_max_abs(const test_local_golden_case *tc,
|
||||
const float *cand_logits,
|
||||
int vocab,
|
||||
int n) {
|
||||
float max_abs = 0.0f;
|
||||
if (n > tc->ntop) n = tc->ntop;
|
||||
for (int i = 0; i < n; i++) {
|
||||
const int id = tc->top[i].id;
|
||||
if (id < 0) continue;
|
||||
if (id < 0 || id >= vocab) return FLT_MAX;
|
||||
const float abs_delta = fabsf(cand_logits[id] - tc->top[i].logit);
|
||||
if (abs_delta > max_abs) max_abs = abs_delta;
|
||||
}
|
||||
@@ -6155,7 +6165,7 @@ static void test_local_golden_case_run(ds4_engine *engine,
|
||||
const int top20_overlap = test_local_golden_overlap(tc, cand_top, 20);
|
||||
const int top64_overlap = test_local_golden_overlap(tc, cand_top, 64);
|
||||
const float top20_max_abs =
|
||||
test_local_golden_max_abs(tc, cand_logits, 20);
|
||||
test_local_golden_max_abs(tc, cand_logits, vocab, 20);
|
||||
|
||||
fprintf(stderr,
|
||||
"ds4-test: local golden %s top1 ref=%d cand=%d "
|
||||
@@ -6185,7 +6195,8 @@ static void test_local_golden_case_run(ds4_engine *engine,
|
||||
|
||||
static void test_local_golden_vectors(void) {
|
||||
const char *path = getenv("DS4_TEST_LOCAL_GOLDEN_FILE");
|
||||
if (!path || !path[0]) {
|
||||
const bool default_fixture = !path || !path[0];
|
||||
if (default_fixture) {
|
||||
path = "tests/test-vectors/flash-0731/local-golden.vec";
|
||||
}
|
||||
FILE *fp = fopen(path, "rb");
|
||||
@@ -6212,7 +6223,13 @@ static void test_local_golden_vectors(void) {
|
||||
}
|
||||
|
||||
test_local_golden_case tc;
|
||||
while (test_read_local_golden_case(fp, &tc)) {
|
||||
const bool compatible_fixture = !default_fixture || !ds4_engine_is_glm_dsa(engine);
|
||||
if (!compatible_fixture) {
|
||||
fprintf(stderr, "ds4-test: DeepSeek local golden vectors skipped for %s; "
|
||||
"set DS4_TEST_LOCAL_GOLDEN_FILE to a matching fixture\n",
|
||||
ds4_engine_model_name(engine));
|
||||
}
|
||||
while (compatible_fixture && test_read_local_golden_case(fp, &tc)) {
|
||||
if (!test_fill_local_golden_case(fp, &tc)) break;
|
||||
test_local_golden_case_run(engine, &tc);
|
||||
}
|
||||
@@ -7375,6 +7392,7 @@ static void test_print_help(const char *prog) {
|
||||
puts(" DS4_TEST_LONG_PROMPT=FILE Rendered long-context story fact prompt.");
|
||||
puts(" DS4_TEST_VECTOR_FILE=FILE Official fixture. Default: flash-0731/official.vec.");
|
||||
puts(" DS4_TEST_LOCAL_GOLDEN_FILE=FILE Local fixture. Default: flash-0731/local-golden.vec.");
|
||||
puts(" DeepSeek default fixtures are skipped for GLM; explicit fixtures are always checked.");
|
||||
puts(" DS4_TEST_MPP_EQ_CASE=NAME Run only Tensor equivalence cases whose id contains NAME.");
|
||||
puts(" DS4_TEST_MTP=FILE Legacy MTP support GGUF for --mtp-verify-depth.");
|
||||
puts(" DS4_TEST_DSPARK=FILE DSpark support GGUF for --dspark-verify-depth.");
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
#include "../ds4.c"
|
||||
|
||||
static void require(int ok, const char *what) {
|
||||
if (!ok) {
|
||||
fprintf(stderr, "HC pre repeat FAIL: %s\n", what);
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
require(argc == 3, "expected repeat and reference flags");
|
||||
const bool repeat = atoi(argv[1]) != 0;
|
||||
const bool reference = atoi(argv[2]) != 0;
|
||||
require(setenv("DS4_METAL_DISABLE_GLM53_HC_PRODUCER_FUSE", "1", 1) == 0,
|
||||
"disable compound producer");
|
||||
require(setenv("DS4_METAL_DISABLE_HC_FUSION", reference ? "1" : "0", 1) == 0,
|
||||
"select split/collapse variant");
|
||||
require(setenv("DS4_GLM_DECODE_REPEAT", repeat ? "hc_pre" : "", 1) == 0,
|
||||
"select repeat mode");
|
||||
g_ds4_shape = DS4_SHAPE_GLM53;
|
||||
const uint32_t dim = DS4_N_HC * DS4_N_EMBD;
|
||||
const uint32_t mix = DS4_N_HC * (DS4_N_HC + 2u);
|
||||
const ds4_tensor fn = {.type = DS4_TENSOR_BF16, .abs_offset = 0};
|
||||
const ds4_tensor scale = {.abs_offset = (uint64_t)dim * mix * sizeof(uint16_t)};
|
||||
const ds4_tensor base = {.abs_offset = scale.abs_offset + 3u * sizeof(float)};
|
||||
const ds4_tensor norm = {.abs_offset = base.abs_offset + mix * sizeof(float)};
|
||||
ds4_model model = {0};
|
||||
model.size = (norm.abs_offset + DS4_N_EMBD * sizeof(float) + 16383u) & ~16383ull;
|
||||
model.map = mmap(NULL, model.size, PROT_READ | PROT_WRITE,
|
||||
MAP_PRIVATE | MAP_ANON, -1, 0);
|
||||
require(model.map != MAP_FAILED, "allocate model fixture");
|
||||
uint16_t *weights = (uint16_t *)model.map;
|
||||
for (uint32_t i = 0; i < dim * mix; i++) {
|
||||
const float value = (float)((int)((i * 1664525u + 1013904223u) >> 24) - 128) / 4096.0f;
|
||||
uint32_t bits;
|
||||
memcpy(&bits, &value, sizeof(bits));
|
||||
weights[i] = (uint16_t)(bits >> 16);
|
||||
}
|
||||
for (uint32_t i = 0; i < 3; i++)
|
||||
((float *)(model.map + scale.abs_offset))[i] = 0.5f + 0.25f * i;
|
||||
for (uint32_t i = 0; i < mix; i++)
|
||||
((float *)(model.map + base.abs_offset))[i] = 0.125f * (float)((int)(i % 5) - 2);
|
||||
for (uint32_t i = 0; i < DS4_N_EMBD; i++)
|
||||
((float *)(model.map + norm.abs_offset))[i] = 1.0f + 0.001f * (i % 7);
|
||||
require(ds4_gpu_init(), "initialize Metal");
|
||||
require(ds4_gpu_set_model_map(model.map, model.size), "register model fixture");
|
||||
ds4_glm_gpu_graph graph = {0};
|
||||
const uint32_t counts[] = {dim, mix, mix, DS4_N_EMBD, DS4_N_EMBD};
|
||||
ds4_gpu_tensor *expected[5], *actual[5];
|
||||
for (unsigned i = 0; i < 5; i++) {
|
||||
expected[i] = ds4_gpu_tensor_alloc(counts[i] * sizeof(float));
|
||||
actual[i] = ds4_gpu_tensor_alloc(counts[i] * sizeof(float));
|
||||
require(expected[i] && actual[i], "allocate producer outputs");
|
||||
}
|
||||
graph.hc_flat = actual[0];
|
||||
graph.hc_mix = actual[1];
|
||||
graph.hc_split = actual[2];
|
||||
ds4_gpu_tensor *residual = ds4_gpu_tensor_alloc(dim * sizeof(float));
|
||||
float *input = malloc(dim * sizeof(float));
|
||||
float *a = malloc(dim * sizeof(float)), *b = malloc(dim * sizeof(float));
|
||||
require(residual && input && a && b, "allocate residual and readback");
|
||||
for (unsigned draw = 0; draw < 3; draw++) {
|
||||
for (uint32_t i = 0; i < dim; i++) {
|
||||
input[i] = draw == 0 ? 0.0f :
|
||||
(0.01f * (float)((int)(i % 23) - 11) + 0.002f * (i % 5)) *
|
||||
(draw == 2 && i % 7 == 0 ? 32.0f : 1.0f);
|
||||
}
|
||||
require(ds4_gpu_tensor_write(residual, 0, input, dim * sizeof(float)), "write residual");
|
||||
require(ds4_gpu_begin_commands(), "begin reference");
|
||||
const uint64_t before_reference = ds4_gpu_encoder_count();
|
||||
require(ds4_gpu_rms_norm_plain_tensor(expected[0], residual, dim, DS4_RMS_EPS), "reference RMS");
|
||||
require(ds4_gpu_glm53_matmul_bf16(expected[1], model.map, model.size,
|
||||
fn.abs_offset, dim, mix, expected[0], 1), "reference projection");
|
||||
if (reference) {
|
||||
require(ds4_gpu_hc_split_sinkhorn_tensor(expected[2], expected[1],
|
||||
model.map, model.size, scale.abs_offset, base.abs_offset,
|
||||
DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS), "reference split");
|
||||
require(ds4_gpu_hc_weighted_sum_tensor(expected[3], residual, expected[2],
|
||||
DS4_N_EMBD, DS4_N_HC), "reference collapse");
|
||||
} else {
|
||||
require(ds4_gpu_hc_split_weighted_sum_tensor(expected[3], expected[2],
|
||||
expected[1], residual, model.map, model.size,
|
||||
scale.abs_offset, base.abs_offset, DS4_N_EMBD, DS4_N_HC,
|
||||
DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS), "reference split/collapse");
|
||||
}
|
||||
require(ds4_gpu_rms_norm_weight_tensor(expected[4], expected[3], model.map,
|
||||
model.size, norm.abs_offset, DS4_N_EMBD, DS4_RMS_EPS), "reference weighted RMS");
|
||||
const uint64_t single_dispatches = ds4_gpu_encoder_count() - before_reference;
|
||||
require(ds4_gpu_end_commands(), "finish reference");
|
||||
require(single_dispatches == (reference ? 5u : 4u), "single producer dispatch count");
|
||||
for (unsigned i = 0; i < 5; i++)
|
||||
require(ds4_gpu_tensor_fill_f32(actual[i], -17.0f, counts[i]), "poison outputs");
|
||||
require(ds4_gpu_begin_commands(), "begin graph producer");
|
||||
const uint64_t before = ds4_gpu_encoder_count();
|
||||
require(glm53_graph_hc_pre(&graph, &model, &fn, &scale, &base, &norm,
|
||||
residual, actual[3], actual[4]), "graph producer");
|
||||
const uint64_t dispatches = ds4_gpu_encoder_count() - before;
|
||||
require(ds4_gpu_end_commands(), "finish graph producer");
|
||||
const uint64_t wanted = single_dispatches * (repeat ? 2u : 1u);
|
||||
if (dispatches != wanted) {
|
||||
fprintf(stderr, "HC pre dispatches: got %llu, expected %llu\n",
|
||||
(unsigned long long)dispatches, (unsigned long long)wanted);
|
||||
return 1;
|
||||
}
|
||||
for (unsigned i = 0; i < 5; i++) {
|
||||
const size_t bytes = counts[i] * sizeof(float);
|
||||
require(ds4_gpu_tensor_read(expected[i], 0, a, bytes) &&
|
||||
ds4_gpu_tensor_read(actual[i], 0, b, bytes), "read producer outputs");
|
||||
require(memcmp(a, b, bytes) == 0, "producer output bytes match single execution");
|
||||
}
|
||||
require(ds4_gpu_tensor_read(residual, 0, a, dim * sizeof(float)), "read residual");
|
||||
require(memcmp(a, input, dim * sizeof(float)) == 0, "residual remains unchanged");
|
||||
}
|
||||
for (unsigned i = 0; i < 5; i++) {
|
||||
ds4_gpu_tensor_free(expected[i]);
|
||||
ds4_gpu_tensor_free(actual[i]);
|
||||
}
|
||||
ds4_gpu_tensor_free(residual);
|
||||
free(input); free(a); free(b);
|
||||
ds4_gpu_cleanup();
|
||||
munmap((void *)model.map, model.size);
|
||||
printf("HC pre repeat PASS: repeat=%u reference=%u, three draws, exact bytes and dispatch count\n",
|
||||
repeat, reference);
|
||||
return 0;
|
||||
}
|
||||
+1890
-1
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,51 @@
|
||||
#define _DARWIN_C_SOURCE
|
||||
#include "ds4_gpu.h"
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <sys/mman.h>
|
||||
bool ds4_log_is_tty(FILE *fp){(void)fp;return false;}
|
||||
static void need(int ok,const char *s){if(!ok){fprintf(stderr,"Q8 inputs FAIL %s\n",s);exit(1);}}
|
||||
static uint32_t seed=0x6719bc41;
|
||||
static uint32_t rnd(void){seed^=seed<<13;seed^=seed>>17;seed^=seed<<5;return seed;}
|
||||
int main(void){
|
||||
const unsigned widths[6]={8192,8192,8192,128,128,64};
|
||||
uint64_t offsets[6],size=0;
|
||||
for(unsigned i=0;i<6;i++){offsets[i]=size;size+=(uint64_t)widths[i]*(i<3?128*34:8192);}
|
||||
size=(size+16383u)&~16383ull;
|
||||
uint8_t *model=mmap(NULL,size,PROT_READ|PROT_WRITE,MAP_PRIVATE|MAP_ANON,-1,0);
|
||||
need(model!=MAP_FAILED,"model");
|
||||
for(unsigned s=0;s<6;s++){
|
||||
if(s<3){uint8_t *b=model+offsets[s];for(unsigned k=0;k<widths[s]*128u;k++,b+=34){
|
||||
uint16_t scale=(k%3?0x1800:0x9800);memcpy(b,&scale,2);
|
||||
for(unsigned j=0;j<32;j++)b[2+j]=(uint8_t)(rnd()%255-127);
|
||||
}}else{uint16_t *w=(uint16_t *)(model+offsets[s]);for(unsigned k=0;k<widths[s]*4096u;k++)w[k]=(uint16_t)((rnd()&0x807fu)|0x3c00u);}
|
||||
}
|
||||
need(ds4_gpu_init(),"init");need(ds4_gpu_set_model_map(model,size),"register");
|
||||
ds4_gpu_tensor *out[6],*x=ds4_gpu_tensor_alloc(4096*4);
|
||||
float *expected[6],*observed[6];
|
||||
for(unsigned i=0;i<6;i++){out[i]=ds4_gpu_tensor_alloc(widths[i]*4);expected[i]=malloc(widths[i]*4);observed[i]=malloc(widths[i]*4);need(out[i]&&expected[i]&&observed[i],"outputs");}
|
||||
need(x!=NULL,"input");setenv("DS4_GLM_ENABLE_KDA_Q8_INPUTS","1",1);
|
||||
for(unsigned draw=0;draw<120;draw++){
|
||||
float input[4096];for(unsigned k=0;k<4096;k++)input[k]=draw%10?((int)(rnd()%4095)-2047)/1024.0f:0;
|
||||
need(ds4_gpu_tensor_write(x,0,input,sizeof(input)),"input write");
|
||||
for(unsigned i=0;i<6;i++){
|
||||
int ok=i<3?ds4_gpu_matmul_q8_0_tensor(out[i],model,size,offsets[i],4096,widths[i],x,1):ds4_gpu_glm53_matmul_bf16(out[i],model,size,offsets[i],4096,widths[i],x,1);
|
||||
need(ok,"reference");need(ds4_gpu_tensor_read(out[i],0,expected[i],widths[i]*4),"reference read");
|
||||
memset(observed[i],0xa5,widths[i]*4);need(ds4_gpu_tensor_write(out[i],0,observed[i],widths[i]*4),"poison");
|
||||
}
|
||||
need(ds4_gpu_glm53_kda_inputs_q8_bf16(out,offsets,model,size,x),"fused");
|
||||
for(unsigned i=0;i<6;i++){
|
||||
need(ds4_gpu_tensor_read(out[i],0,observed[i],widths[i]*4),"fused read");
|
||||
if(memcmp(expected[i],observed[i],widths[i]*4)){unsigned k=0;while(k<widths[i]&&!memcmp(expected[i]+k,observed[i]+k,4))k++;fprintf(stderr,"draw=%u slot=%u row=%u ref=%a observed=%a\n",draw,i,k,expected[i][k],observed[i][k]);return 1;}
|
||||
}
|
||||
}
|
||||
uint64_t saved=offsets[5];offsets[5]=size-1;
|
||||
need(!ds4_gpu_glm53_kda_inputs_q8_bf16(out,offsets,model,size,x),"invalid weight refusal");offsets[5]=saved;
|
||||
unsetenv("DS4_GLM_ENABLE_KDA_Q8_INPUTS");
|
||||
need(!ds4_gpu_glm53_kda_inputs_q8_bf16(out,offsets,model,size,x),"disabled refusal");
|
||||
for(unsigned i=0;i<6;i++){ds4_gpu_tensor_free(out[i]);free(expected[i]);free(observed[i]);}
|
||||
ds4_gpu_tensor_free(x);ds4_gpu_cleanup();munmap(model,size);
|
||||
puts("GLM mixed Q8/BF16 inputs exact PASS: 120 poisoned draws, six complete projections, refusal gates");return 0;
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
#define _DARWIN_C_SOURCE
|
||||
#include "ds4_gpu.h"
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <sys/mman.h>
|
||||
|
||||
bool ds4_log_is_tty(FILE *fp) { (void)fp; return false; }
|
||||
|
||||
static void require(int ok, const char *what) {
|
||||
if (!ok) { fprintf(stderr, "router/shared FAIL: %s\n", what); exit(1); }
|
||||
}
|
||||
static uint32_t state = 0x514aac39u;
|
||||
static uint32_t random_bits(void) {
|
||||
state ^= state << 13; state ^= state >> 17; state ^= state << 5; return state;
|
||||
}
|
||||
static float sample(void) { return ((int)(random_bits() & 2047u)-1024) / 1024.0f; }
|
||||
|
||||
int main(void) {
|
||||
enum { WIDTH=4096, EXPERTS=288, MID=2048, USED=8,
|
||||
ROUTER=0, BIAS=WIDTH*EXPERTS*4,
|
||||
GATE=BIAS+16384, UP=GATE+128*34*MID,
|
||||
BYTES=UP+128*34*MID+16384 };
|
||||
uint8_t *model=mmap(NULL,BYTES,PROT_READ|PROT_WRITE,MAP_PRIVATE|MAP_ANON,-1,0);
|
||||
require(model!=MAP_FAILED,"model allocation");
|
||||
for (unsigned i=0;i<WIDTH*EXPERTS;i++) ((float *)model)[i]=sample()/32.0f;
|
||||
for (int arm=0;arm<2;arm++) {
|
||||
uint8_t *p=model+(arm?UP:GATE);
|
||||
for (unsigned b=0;b<128u*MID;b++,p+=34) {
|
||||
const uint16_t scale=(b%7==0 ? 0xa000u : 0x2000u);
|
||||
memcpy(p,&scale,2);
|
||||
for (unsigned q=0;q<32;q++) p[2+q]=(uint8_t)(random_bits()%255u-127u);
|
||||
}
|
||||
}
|
||||
require(ds4_gpu_init(),"GPU initialization");
|
||||
require(ds4_gpu_set_model_map(model,BYTES),"model registration");
|
||||
const uint64_t sizes[]={EXPERTS*4,USED*4,USED*4,EXPERTS*4,4,MID*4};
|
||||
ds4_gpu_tensor *out[6];
|
||||
uint8_t *expected[6], *observed[6];
|
||||
for (int i=0;i<6;i++) {
|
||||
out[i]=ds4_gpu_tensor_alloc(sizes[i]);
|
||||
expected[i]=malloc(sizes[i]); observed[i]=malloc(sizes[i]);
|
||||
require(out[i] && expected[i] && observed[i],"output allocation");
|
||||
}
|
||||
ds4_gpu_tensor *x=ds4_gpu_tensor_alloc(WIDTH*4);
|
||||
require(x!=NULL,"input allocation");
|
||||
require(ds4_gpu_tensor_fill_f32(out[4],0.0f,1u),"counter reset");
|
||||
for (unsigned draw=0;draw<240;draw++) {
|
||||
float input[WIDTH];
|
||||
for (unsigned i=0;i<WIDTH;i++) input[i]=(draw%12==0)?0.0f:sample();
|
||||
for (unsigned i=0;i<EXPERTS;i++) ((float *)(model+BIAS))[i]=(draw%12==0)?0.0f:sample()/4.0f;
|
||||
if (draw%12==1) {
|
||||
uint32_t nan=0x7fc00129; memcpy(model+BIAS+4*179,&nan,4);
|
||||
}
|
||||
if (draw%12==2) {
|
||||
uint32_t inf=0x7f800000; memcpy(model+BIAS+4*219,&inf,4);
|
||||
}
|
||||
require(ds4_gpu_tensor_write(x,0,input,sizeof(input)),"input write");
|
||||
require(ds4_gpu_matmul_f32_tensor(out[0],model,BYTES,ROUTER,WIDTH,EXPERTS,x,1),"reference projection");
|
||||
require(ds4_gpu_glm_router_select_tensor(out[1],out[2],out[3],model,BYTES,BIAS,out[0],EXPERTS,USED,1.0f),"reference selection");
|
||||
require(ds4_gpu_shared_mid_swiglu_q8_0_tensor(out[5],model,BYTES,GATE,UP,WIDTH,MID,x,10.0f),"reference shared mid");
|
||||
for (int i=0;i<6;i++) {
|
||||
require(ds4_gpu_tensor_read(out[i],0,expected[i],sizes[i]),"reference read");
|
||||
if (i!=4) {
|
||||
memset(observed[i],0xa5,sizes[i]);
|
||||
require(ds4_gpu_tensor_write(out[i],0,observed[i],sizes[i]),"output poison");
|
||||
}
|
||||
}
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,1.0f,10.0f)==1,"fused dispatch");
|
||||
for (int i=0;i<6;i++) {
|
||||
require(ds4_gpu_tensor_read(out[i],0,observed[i],sizes[i]),"fused read");
|
||||
if (memcmp(expected[i],observed[i],sizes[i])) {
|
||||
unsigned bad=0;while(bad<sizes[i]/4 && ((uint32_t *)expected[i])[bad]==((uint32_t *)observed[i])[bad])bad++;
|
||||
fprintf(stderr,"draw=%u output=%d word=%u expected=%08x got=%08x\n",draw,i,bad,((uint32_t *)expected[i])[bad],((uint32_t *)observed[i])[bad]);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
uint32_t partial_arrivals=100u;
|
||||
require(ds4_gpu_tensor_write(out[4],0,&partial_arrivals,sizeof(partial_arrivals)),
|
||||
"poison partial router arrivals");
|
||||
ds4_gpu_test_invalidate_completion_counters();
|
||||
for (int i=0;i<6;i++) {
|
||||
if (i!=4) {
|
||||
memset(observed[i],0xa5,sizes[i]);
|
||||
require(ds4_gpu_tensor_write(out[i],0,observed[i],sizes[i]),"recovery output poison");
|
||||
}
|
||||
}
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,1.0f,10.0f)==1,"recovered fused dispatch");
|
||||
for (int i=0;i<6;i++) {
|
||||
require(ds4_gpu_tensor_read(out[i],0,observed[i],sizes[i]),"recovered fused read");
|
||||
require(memcmp(expected[i],observed[i],sizes[i])==0,"failure invalidation restores complete outputs");
|
||||
}
|
||||
const char *disabled[] = {"DS4_METAL_DISABLE_GLM53_FLASH_TUNING",
|
||||
"DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE", "DS4_METAL_DISABLE_GLM53_ROUTER_TOP8"};
|
||||
for (unsigned i=0;i<sizeof(disabled)/sizeof(disabled[0]);i++) {
|
||||
setenv(disabled[i], "1", 1);
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,1.0f,10.0f)==0,"aggregate rollback refusal");
|
||||
unsetenv(disabled[i]);
|
||||
}
|
||||
const uint32_t nonfinite_bits[]={0x7fc00123u,0x7f800000u,0xff800000u};
|
||||
for (unsigned i=0;i<sizeof(nonfinite_bits)/sizeof(nonfinite_bits[0]);i++) {
|
||||
float nonfinite;
|
||||
memcpy(&nonfinite,&nonfinite_bits[i],sizeof(nonfinite));
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,nonfinite,10.0f)==0,"nonfinite scale refusal");
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,1.0f,nonfinite)==0,"nonfinite clamp refusal");
|
||||
}
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,BYTES-1,x,1.0f,10.0f)==0,"late model range refusal");
|
||||
ds4_gpu_set_ssd_streaming(true);
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,1.0f,10.0f)==0,"SSD refusal");
|
||||
ds4_gpu_set_ssd_streaming(false);
|
||||
require(setenv("DS4_METAL_DISABLE_GLM53_ROUTER_SHARED","1",1)==0,"set rollback");
|
||||
require(ds4_gpu_glm53_router_shared_exact(out[0],out[1],out[2],out[3],out[4],out[5],
|
||||
model,BYTES,ROUTER,BIAS,GATE,UP,x,1.0f,10.0f)==0,"rollback refusal");
|
||||
for (int i=0;i<6;i++) { ds4_gpu_tensor_free(out[i]);free(expected[i]);free(observed[i]); }
|
||||
ds4_gpu_tensor_free(x);ds4_gpu_cleanup();munmap(model,BYTES);
|
||||
puts("GLM router/shared exact: PASS (240 poisoned draws, failure recovery, counter reuse)");
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
#define _DARWIN_C_SOURCE
|
||||
#include "ds4_gpu.h"
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <time.h>
|
||||
bool ds4_log_is_tty(FILE *fp) { (void)fp; return false; }
|
||||
static void need(int ok,const char *why) { if(!ok){fprintf(stderr,"TOPK FAIL %s\n",why);exit(1);} }
|
||||
static uint32_t seed=0xb719325e;
|
||||
static uint32_t rnd(void){seed^=seed<<13;seed^=seed>>17;seed^=seed<<5;return seed;}
|
||||
static float bits(uint32_t b){float f;memcpy(&f,&b,4);return f;}
|
||||
int main(void) {
|
||||
enum { MAX=131073, K=512 };
|
||||
const uint32_t widths[]={8192,12288,15543,37500,75000,131073};
|
||||
float *x=malloc(MAX*4); uint32_t a[K],b[K],stats[15]={0},before[15]={0};
|
||||
need(x!=NULL,"CPU allocation"); need(ds4_gpu_init(),"init");
|
||||
ds4_gpu_tensor *scores=ds4_gpu_tensor_alloc(MAX*4),*out=ds4_gpu_tensor_alloc(K*4);
|
||||
need(scores&&out,"GPU allocation");
|
||||
unsigned accepts=0,rejects=0,skips=0;
|
||||
for(unsigned w=0;w<sizeof(widths)/sizeof(widths[0]);w++)
|
||||
for(unsigned draw=0;draw<120;draw++) {
|
||||
const uint32_t n=widths[w], kind=draw%12;
|
||||
for(uint32_t i=0;i<n;i++) x[i]=(float)(i+1);
|
||||
if(kind==1) for(uint32_t i=0;i<n;i++) x[i]=-x[i];
|
||||
if(kind==2) for(uint32_t i=0;i<n;i++) x[i]-=(float)n/2.0f;
|
||||
if(kind==3) x[n-2]=x[n-1]; // tie among winners
|
||||
if(kind==4) x[n-K-1]=x[n-K]; // tie at cut
|
||||
if(kind==5) x[0]=bits(0x7fc01234); // nonfinite outside winners
|
||||
if(kind==6) x[0]=bits(0x7f800000);
|
||||
if(kind==7) x[0]=bits(1); // subnormal outside winners
|
||||
if(kind==8) for(uint32_t i=0;i<n;i++) x[i]=bits(0x3f800000+i); // overflow
|
||||
if(kind==9) for(uint32_t i=0;i<n;i++) x[i]=bits(i&1?0x80000000:0);
|
||||
if(kind==10) for(uint32_t i=0;i<n;i++) x[i]=bits(0x00800000+(i<<6));
|
||||
if(kind==11) for(uint32_t i=0;i<n;i++) x[i]=bits(0xff7fffff-(i<<6));
|
||||
for(uint32_t i=n-1;i;i--){uint32_t j=rnd()%(i+1);float t=x[i];x[i]=x[j];x[j]=t;}
|
||||
need(ds4_gpu_tensor_write(scores,0,x,n*4),"write");
|
||||
unsetenv("DS4_GLM_ENABLE_TOPK_FAST");
|
||||
need(ds4_gpu_indexer_topk_tensor(out,scores,n,1,K),"reference");
|
||||
need(ds4_gpu_tensor_read(out,0,a,sizeof(a)),"read reference");
|
||||
memcpy(before,stats,sizeof(stats));
|
||||
setenv("DS4_GLM_ENABLE_TOPK_FAST","1",1);
|
||||
memset(b,0xa5,sizeof(b));need(ds4_gpu_tensor_write(out,0,b,sizeof(b)),"poison");
|
||||
need(ds4_gpu_glm53_indexer_topk_tensor(out,scores,n,1,K),"candidate");
|
||||
need(ds4_gpu_tensor_read(out,0,b,sizeof(b)),"read candidate");
|
||||
if(memcmp(a,b,sizeof(a))){unsigned i=0;while(i<K&&a[i]==b[i])i++;fprintf(stderr,"n=%u draw=%u kind=%u rank=%u ref=%u got=%u\n",n,draw,kind,i,a[i],b[i]);return 1;}
|
||||
ds4_gpu_test_glm53_topk_stats(stats,15);
|
||||
if(stats[4]==before[4]){skips++;need(n==8192||n==131073,"unexpected skip");}
|
||||
else {need(stats[4]==before[4]+1,"call count");if(stats[5]>before[5]){accepts++;need(kind==0||kind==1||kind==2||kind==10||kind==11,"unsafe accept");}else rejects++;}
|
||||
}
|
||||
const uint32_t recovery_n=12288u;
|
||||
for(uint32_t i=0;i<recovery_n;i++) x[i]=-1.0f;
|
||||
for(uint32_t i=0;i<K;i++) x[recovery_n-K+i]=bits(0x3f800000u+i);
|
||||
need(ds4_gpu_tensor_write(scores,0,x,recovery_n*4),"recovery write");
|
||||
need(ds4_gpu_indexer_topk_tensor(out,scores,recovery_n,1,K),"recovery reference");
|
||||
need(ds4_gpu_tensor_read(out,0,a,sizeof(a)),"recovery reference read");
|
||||
need(ds4_gpu_test_glm53_topk_poison_recovery_state(),"poison stale 512 candidates");
|
||||
need(ds4_gpu_glm53_indexer_topk_tensor(out,scores,recovery_n,1,K),"poisoned selector");
|
||||
need(ds4_gpu_tensor_read(out,0,b,sizeof(b)),"poisoned selector read");
|
||||
need(memcmp(a,b,sizeof(a))!=0,"stale 512 candidates reproduce contamination");
|
||||
need(ds4_gpu_test_glm53_topk_poison_recovery_state(),"repoison stale 512 candidates");
|
||||
ds4_gpu_test_invalidate_completion_counters();
|
||||
memset(b,0xa5,sizeof(b));need(ds4_gpu_tensor_write(out,0,b,sizeof(b)),"recovery output poison");
|
||||
need(ds4_gpu_glm53_indexer_topk_tensor(out,scores,recovery_n,1,K),"recovered selector");
|
||||
need(ds4_gpu_tensor_read(out,0,b,sizeof(b)),"recovered selector read");
|
||||
need(memcmp(a,b,sizeof(a))==0,"failure invalidation drops stale selector state");
|
||||
need(accepts>80&&rejects>160&&skips==240,"non-vacuous acceptance/fallback/host gates");
|
||||
printf("GLM DSA top-k exact PASS: 720 poisoned draws, recovery, accepted=%u fallback=%u host-skipped=%u\n",accepts,rejects,skips);
|
||||
ds4_gpu_tensor_free(scores);ds4_gpu_tensor_free(out);ds4_gpu_cleanup();free(x);return 0;
|
||||
}
|
||||
@@ -433,6 +433,46 @@ static void test_glm_attention_budget(void) {
|
||||
assert(glm_graph_dense_compact_attention_limit(g) == 2051);
|
||||
g->full_kv_cache = false;
|
||||
}
|
||||
|
||||
const char *resident_env = getenv("DS4_GLM_FULL_ATTN_CAP");
|
||||
const char *streaming_env = getenv("DS4_GLM_FULL_ATTN_STREAMING_CAP");
|
||||
char *saved_resident = resident_env ? strdup(resident_env) : NULL;
|
||||
char *saved_streaming = streaming_env ? strdup(streaming_env) : NULL;
|
||||
assert(!resident_env || saved_resident);
|
||||
assert(!streaming_env || saved_streaming);
|
||||
assert(setenv("DS4_GLM_FULL_ATTN_CAP", "1024", 1) == 0);
|
||||
assert(setenv("DS4_GLM_FULL_ATTN_STREAMING_CAP", "1024", 1) == 0);
|
||||
assert(glm_graph_full_attention_cap(4096, false) == 1024);
|
||||
assert(glm_graph_full_attention_cap(4096, true) == 1024);
|
||||
const uint32_t selected_capacity =
|
||||
glm_graph_indexer_selected_capacity(4096, false);
|
||||
assert(selected_capacity == 2051);
|
||||
assert(glm_graph_indexer_selected_capacity(1024, false) == 1024);
|
||||
assert(ds4_gpu_init());
|
||||
ds4_gpu_tensor *selected = ds4_gpu_tensor_alloc(
|
||||
(uint64_t)selected_capacity * sizeof(uint32_t));
|
||||
assert(selected);
|
||||
assert(ds4_gpu_glm_fill_selected_range_tensor(selected, 1025));
|
||||
uint32_t last = UINT32_MAX;
|
||||
assert(ds4_gpu_tensor_read(selected, 1024u * sizeof(uint32_t),
|
||||
&last, sizeof(last)));
|
||||
assert(last == 1024);
|
||||
printf("GLM attention cap: configured=1024 selected_capacity=%u "
|
||||
"filled=1025 last=%u: PASS\n", selected_capacity, last);
|
||||
ds4_gpu_tensor_free(selected);
|
||||
ds4_gpu_cleanup();
|
||||
if (saved_resident) {
|
||||
assert(setenv("DS4_GLM_FULL_ATTN_CAP", saved_resident, 1) == 0);
|
||||
} else {
|
||||
assert(unsetenv("DS4_GLM_FULL_ATTN_CAP") == 0);
|
||||
}
|
||||
if (saved_streaming) {
|
||||
assert(setenv("DS4_GLM_FULL_ATTN_STREAMING_CAP", saved_streaming, 1) == 0);
|
||||
} else {
|
||||
assert(unsetenv("DS4_GLM_FULL_ATTN_STREAMING_CAP") == 0);
|
||||
}
|
||||
free(saved_streaming);
|
||||
free(saved_resident);
|
||||
free(g);
|
||||
g_ds4_shape = saved_shape;
|
||||
}
|
||||
@@ -519,6 +559,8 @@ static void test_glm_spec_rollback(void) {
|
||||
for (int j = 0; j < 128; j++) assert(state[j] == 11);
|
||||
}
|
||||
}
|
||||
puts("GLM speculative rollback: indexer tail/key cache and KDA conv/"
|
||||
"recurrent state restored across replay: PASS");
|
||||
ds4_gpu_tensor_free(x);
|
||||
ds4_gpu_tensor_free(gate);
|
||||
glm_graph_free(g);
|
||||
|
||||
Reference in new issue
Block a user