Module ops

Expand description

Shared kernel dispatch operations.

Freestanding functions wrapping CUDA kernel launches via KernelLaunch. Layer implementations compose these to build forward passes.

Each function’s parameters exactly match the corresponding CUDA kernel signature. Grid/block dimensions are computed from the problem size.

Refactor wave 4a (2026-05-03): split into ops/ sub-modules with thematic groupings. All public functions remain available at this path via re-export.

Re-exports§

pub use model_stats::ModelStats;
pub use moe_lora_grouped::*;

Modules§

gdn_flashinfer: Opt-in FlashInfer GDN prefill via dlopen(libatlasgdn.so) — behind ATLAS_GDN_FLASHINFER=1.
lora_delta: Runtime LoRA delta: y += scale * (x @ A^T) @ B^T, BF16 side-path. Zero new CUDA kernels — reuses dense_gemv_bf16 / dense_gemm_tc / dense_gemm_bf16 / bf16_scaled_add, all shipped in kernels/gb10/common/.
model_stats: ModelStats — diagnostic counters and one-shot latches owned by the model.
moe_lora_grouped: Device-side MoE expert down_proj LoRA fold launcher (moe_lora_grouped_down).
token_overlay: Token-overlay kernel launchers (Feature 2). Thin KernelLaunch wrappers over kernels/gb10/common/token_overlay.cu:

Structs§

DenseMmKernels: The three BF16 dense kernels one projection site can land on, resolved once.
DerivedWeights: Per-model memo of derived weight encodings.
GemmDispatch: Which GEMM implementation each projection takes.
Glm5NextMhcKernels: Every kernel GLM-5.3’s hyper-connection needs, all from the single module glm5next_mhc.
Glm5NextMhcSiteWeights: Per-site mHC weights. One set for the attention site, one for the FFN site.
ModelLevers: Kernel-path levers for one loaded model.
MoeCutlassDownHostTables: Down-projection third of MoeCutlassHostTables.
MoeCutlassHostTables: Host snapshots of the per-expert pointer/scale tables for the CUTLASS grouped path, owned by the MoeLayer whose device tables they mirror.

Enums§

Derivation: Which derivation a lookup is for.
HcVariant: Which family a site’s weights select.

Constants§

DENSE_GEMV_BATCHM_DECODE_MAX_M: The band the batched GEMV is allowed to CLAIM on the decode paths: the MTP row dispatch and the BF16 lm_head arm.
DENSE_GEMV_BATCHM_MAX_M: Dense BF16 batched GEMV (M rows): C[t] = A[t] @ B^T for t in [0, M).
GLM5NEXT_MHC_MODULE: The one module name GLM’s mHC resolves from.
MAX_DYNAMIC_SMEM: Largest dynamic shared-memory block a kernel may opt into on the GB10 target. Measured on the device (sm_121): CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 101376 (static default is only 49152; per-SM total is 102400).
MHC_MIX_MAX_TOKENS: Token bound on the mix scratch. The GLM stack drives mHC one token at a time (the highway forces a serial prefill), so this is slack, not a shape — but glm_hc_pre REFUSES above it rather than writing past the allocation.
MOE_TOPK_SIGMOID_MAX_EXPERTS: Largest num_experts the sigmoid routing kernels can hold, from #define MAX_EXPERTS in the same file. Beyond it the kernel silently considers only the first MAX_EXPERTS experts (actual_n is a min), so routing stays memory-safe but stops matching the checkpoint.
MOE_TOPK_SIGMOID_MAX_TOP_K: Largest top_k the sigmoid routing kernels can hold.
NVFP4_BLOCK_BYTES
NVFP4_MMQ_SMEM: Dynamic shared memory: ids(512) + x-tile(128MMQ_MMA_TILE_X_K_FP4=764) + y-tile(128*144).
Q2_0_BLOCK_BYTES: sizeof(block_q2_0) bytes: fp16 scale d (2) + 128 codes @ 4/byte (32) = 34.
Q2_BATCHM_MAX_M: Shared-memory row cap of q2_0_gemv_vec_batchm (MAX_M in the .cu). The kernel stages exactly M activation rows in s_A[MAX_M * TILE_K]; passing M > MAX_M overflows that tile (OOB smem write) AND drops output rows >= 8 (compute/write loops iterate m < MAX_M). Callers with more rows MUST chunk — done transparently by q2_0_gemv_vec_batchm.
Q4K_BLOCK_BYTES: sizeof(block_q4_K) bytes.
Q4K_MMQ_SMEM: Dynamic shared memory for the Q4_K MMQ kernel (mmq_x=mmq_y=128, GB10). >48KB -> registry sets attr.
QK2_0: Q2_0 MMQ block size: 128 weights per block_q2_0.
QK_K: Q4_K block size (256 weights -> 144-byte block_q4_K).
QK_NVFP4: NVFP4 block: 64 weights -> 36-byte block_nvfp4 {4×ue4m3 scales, 32B e2m1 nibbles}.
SSD_L: SSD chunk length (must match SSD_L in the kernel).
SSD_PT: head_dim rows per SSD scan block (must match SSD_PT in the kernel).
W4A16_GEMV_OUTS_PER_BLOCK: Base w4a16_gemv: 4 outputs / 256-thread block. SSOT with kernels/**/w4a16_gemv.cu #define N_PER_BLOCK 4.
W4A16_GEMV_SW_OUTS_PER_BLOCK: Single-warp w4a16_gemv_sw: 8 outputs / 256-thread block. SSOT with #define N_PER_BLOCK_SW 8.

Functions§

argmax_bf16: GPU-side argmax over BF16 logits.
argmax_bf16_batch: Batched argmax: ONE launch, one block per row, instead of n serial launches of the single-row argmax_bf16 (which is a one-CTA reduction and so uses 1 of 48 SMs). Byte-identical — each block runs the identical per-row body.
argmax_bf16_batch_lp: Batched argmax that ALSO writes each row’s top-1 log-probability (out_logprob[row] = log softmax(row)[argmax], FP32), computed by online softmax in the same pass — same bandwidth as argmax_bf16_batch, same index semantics.
batched_embed: Batched embedding: gather N rows from embedding table in one launch.
batched_embed_fp8: FP8-table variant of batched_embed: rows are FP8 E4M3 bytes with a per-row f32 dequant scale (the quantize_bf16_to_fp8 layout); the kernel dequantizes on read and writes BF16 rows.
bf16_absmax: Compute max absolute value of a BF16 buffer into a device-side f32.
bf16_concat: BF16 concatenation: out[0..N] = a[0..N], out[N..2N] = b[0..N].
bf16_to_fp8: Convert BF16 activations to FP8 E4M3 for FP8×FP8 GEMM.
compute_gdn_gates: Compute GDN gates from interleaved BA projection + learned A_log/dt_bias.
conv1d_fwd: Full-sequence causal depthwise conv1d + SiLU activation.
conv1d_update: Causal conv1d update (decode step, supports batched sequences).
conv1d_update_chunk2: Fused 2-token conv1d sliding window update + SiLU.
conv1d_update_l2norm: Fused conv1d update + SiLU + L2 normalization for Q/K channels.
conv1d_update_l2norm_strided: conv1d_update_l2norm with INDEPENDENT input/output row strides, so N concurrent decode sequences go in ONE launch instead of N.
conv1d_update_prefill: Multi-token conv1d sliding window update + SiLU for prefill.
cublas_bf16_proj: Route a projection out[M,N] = act[M,K] @ weightᵀ through cuBLASLt BF16. The FP8 weight is dequantized to BF16 once (cached); W16A16 here is strictly more accurate than the blockscaled W8A8 path it replaces.
cublas_bf16_proj_dense: Route a projection out[M,N] = act[M,K] @ weightᵀ through cuBLASLt BF16 for a weight that is already native BF16 [N,K] (no dequant step). Used by models whose attention/shared-expert weights ship unquantized (e.g. Laguna), which can never satisfy the as_fp8() gate of cublas_bf16_proj.
cublas_fp8_proj: Route a projection through native-FP8 cuBLASLt block-scaled matmul: quantize the activation to FP8 + per-[token,128-of-K] VEC128 scales (the existing per_token_group_quant_fp8 kernel), feed the FP8 weight + its per-128×128 block scales directly (zero dequant, zero extra weight memory). Both operands 128-block-scaled (cuBLASLt requires it). ~1.8× the bf16 path (152 vs 85 TF).
cublas_fp8_rowwise_proj: Route a projection through ROW-WISE native-FP8 cuBLASLt (the fp8 path GB10 supports). Weight is re-quantized once to per-row fp8 (cached); the activation is quantized per-token each call. ~1.8× the bf16 path (152 vs 85 TF), and frees the bf16-dequant memory the bf16 path holds. act_fp8_scratch ≥ m*k fp8 bytes; act_scale_scratch ≥ m f32 (e.g. the buffers.fp8_act / fp8_act_scale arena buffers).
cutlass_bf16_proj: Route a projection out[M,N] = act[M,K] @ weightᵀ through CUTLASS BF16.
cutlass_nvfp4_proj
cutlass_nvfp4_proj_from_fp8: Native CUTLASS NVFP4 projection for FP8 checkpoint weights. The FP8 weight is dequantized to BF16 using the existing cache, then packed once into Atlas-transposed NVFP4 data/scales and reused for future calls.
deinterleave_qg: Deinterleave Q/Gate from per-head interleaved to contiguous layout (in-place).
deinterleave_qg_split: Deinterleave Q/Gate with split output — Q to separate contiguous buffer.
deinterleave_qg_split_qnorm: Fused deinterleave Q/Gate + per-head Q RMS norm.
deinterleave_qg_split_qnorm_mrope: Fused deinterleave Q/Gate + per-head Q RMS norm + MRoPE.
deinterleave_qkvz: Deinterleave QKVZ projection output from GQA-grouped to sequential layout.
dense_gemm
dense_gemm_ba_gates_prefill: Fused BA GEMM + GDN gates for prefill (token-parallel).
dense_gemm_bf16_pipelined: Pipelined tensor-core BF16 GEMM — drop-in faster dense_gemm (kernel dense_gemm_bf16_pipelined): mma.sync.m16n8k16 + cp.async 2-stage, 128x128 tile. ~40x the scalar dense_gemm on large-M shapes (cosine=1.0, same math). Grid: (ceil(N/128), ceil(M/128), 1) Block: (256, 1, 1)
dense_gemm_prefill: Dense BF16 prefill GEMM. Prefer the pipelined tensor-core kernel when the selected target ships it, and retain the scalar kernel as an explicit compatibility fallback for older targets.
dense_gemm_router: Order-preserving register-blocked BF16 GEMM (kernel dense_gemm_bf16_router).
dense_gemm_splitk: Split-K GEMM: partial products over K_splits chunks, then reduce. Uses FP32 workspace of size K_splits * M * N * 4 bytes.
dense_gemm_tc: Dense BF16 GEMM: C = A @ B^T.
dense_gemm_tc_scaled_acc: output[m, n] += scale * bf16(input[m, k] @ weight[n, k]^T) in ONE pass.
dense_gemv: Dense BF16 GEMV (M=1): C = A @ B^T for single-row activations.
dense_gemv_ba_gates: Fused BA projection + GDN gates: dense GEMV + gate/beta transforms.
dense_gemv_batch2: Dense BF16 GEMV, batched over 2 rows (M=2): one pass over the weight produces both output rows, halving weight bandwidth vs two dense_gemv launches. Bit-identical to two M=1 dense_gemv calls — each row’s accumulator follows the same K-iteration/reduction order.
dense_gemv_batchm
dense_gemv_fp8w: Dense FP8-weight GEMV (M=1): C = A @ (dequant(B_fp8) * row_scale).
dense_gemv_fp8w_batch2: FP8-weight dual-GEMV. input is [2, K] BF16, output is [2, N] BF16. Grid: (ceil(N/4), 1, 1) Block: (256, 1, 1)
dense_mm_bf16: C[M, N] = A[M, K] @ B[N, K]^T, BF16 in and out, output row stride N.
dequant_nvfp4_to_bf16: Dequantize NVFP4 weight [n, k] (packed E2M1 + E4M3 group scales + per-tensor scale2) -> bf16 [n, k].
dequant_q2_0_gn_to_bf16: Dequant a packed Q2_0 weight [N, K] (contiguous block_q2_0 blocks) into a pre-allocated BF16 scratch buffer [N, K] on stream, IN PLACE (no alloc, no host sync). Reuses the load-time dequant_q2_0_gn_to_bf16 kernel (dequant_gguf_bf16 module). Used by packed-Q2 PREFILL: dequant → transient BF16 → normal BF16 GEMM → free scratch (the resident weight stays 2-bit).
embed_from_argmax: GPU-side argmax + embedding lookup — eliminates D2H sync in MTP propose.
fill_slots_from_block_table: Fill paged-KV slot mappings on-device from a persistent block table.
fp4_act_scratch_bytes: block_fp4_mmq activation scratch bytes for [m, k]. +1MB slack: the kernel’s smem copy loop rounds the last y-slice read up to warp granularity (same convention as q8_1_scratch_bytes). Always ≤ q8_1_scratch_bytes(m, k) → fits the shared ffn_act_q8.
fp8_fp8_gemm_m128_mfast: fp8_fp8_gemm_t_m128_mfast: FP8 A x FP8 B, 128-row M tile, m on the fast axis. A must already be E4M3 (see bf16_to_fp8); the MMA consumed E4M3 either way, so pre-casting A is numerically identical to the BF16-A kernel.
fp8_fp8_gemm_n128: FP8×FP8 GEMM: A [M, K] FP8 × B [N, K] FP8 → C [M, N] BF16.
fp8_fp8_gemm_n128_m128: M128 variant of fp8_fp8_gemm_n128: halves B re-reads for large M (ISL > 128).
fp8_gemm_m128_mfast: fp8_gemm_t_m128_mfast: 128-row M tile (2 chunks/CTA), m on the fast axis. Halves the B panel passes relative to fp8_gemm_n128_mfast.
fp8_gemm_n128: Pre-dequanted FP8 GEMM (prefill): C = A @ B_fp8.
fp8_gemm_n128_m128: M128 variant of fp8_gemm_n128: halves B re-reads for large M (ISL > 128).
fp8_gemm_n128_mfast: Pre-dequant NVFP4 → FP8 E4M3. One-time conversion at model load.
fp8_gemm_n128_row_scaled: Row-scaled FP8 GEMM: C[M, N] = A[M, K] @ (dequant(B_fp8[N, K]) * row_scale[N]).
fp8_gemm_n128_row_scaled_m16: Small-M row-scaled FP8 GEMM (M ≤ 16) — single warp per CTA variant.
fp8_gemm_t_blockscaled: W8A8 + FP32 epilogue GEMM with per-token activation scales and per-block weight scales — vLLM-equivalent FP8 numerics.
fp8_gemv_rowscale_batch8_rt2: Register-tiled batched row-scaled FP8 GEMV (M<=8, T=2 outputs/thread) — the FP8 twin of w4a16_gemv_batch8_rt2, for the DFlash drafter PROPOSE path. input [M, K] BF16, output [M, N] BF16; per-row f32 scale applied at write-out inside the kernel. Replaces the prefill-class tile GEMMs (fp8_gemm_t_row_scaled M64-tile / _m16) that pad 87%/50% of their M-tile at M=8 (~100 GB/s measured vs 180+ for the rt family). Drafter-side numerics: correctness-free under strict-argmax accept. Kernel: fp8_gemv_rowscale_batch8_rt2 (module fp8_gemv_rt). Grid: (ceil(N/8), 1, 1) Block: (256, 1, 1). Requires K % 16 == 0.
fp8_gemv_rowscale_batch16_rt2: MAX_M=16 sibling of fp8_gemv_rowscale_batch8_rt2 for the γ>8 DFlash propose window (flags 9..17). Same template, same launch geometry; added 2026-08-29 after STEP_TIMING measured propose 18.2ms (flag 8, rt2) vs 38.0ms (flag 9, tile fallback) — the whole γ>8 step tax. Kernel: fp8_gemv_rowscale_batch16_rt2 (module fp8_gemv_rt).
fused_k_norm_rope_cache_write_bf16: Fused K-path: rms_norm → RoPE → BF16 paged cache write in one kernel.
fused_k_norm_rope_cache_write_bf16_mrope: MRoPE-interleaved variant — selects abs position from pos_t/pos_h/pos_w based on pair_idx % 3. For text-only inputs (pos_h == pos_w == pos_t) the result is bit-identical to the scalar-position variant.
fused_k_norm_rope_cache_write_fp8: FP8-output sibling of fused_k_norm_rope_cache_write_bf16. Same semantics; one fewer BF16 round before the saturating FP8 cast.
gated_rms_norm: Gated RMS norm (norm_before_gate=False, per-group): output = rms_norm_per_group(input * silu(gate), weight, group_size)
gated_rms_norm_prefill: Batched gated RMS norm for prefill: all (head, actual_token) pairs in one launch.
gdn_decode: Gated delta rule decode (recurrent SSM update, supports batched sequences).
gdn_decode_chunk2: Fused 2-token GDN decode (speculative verification).
gdn_decode_chunk3: Fused 3-token GDN decode (K=3 speculative verification).
gdn_decode_f16_strided_norm: FP16 h-state twin of gdn_decode_f32_strided_norm (ATLAS_SSM_H_FP16).
gdn_decode_f32_conv_norm: Split-v_dim prefill: 2 CTAs per v-head, 64 threads each. FUSED conv1d_update_l2norm + recurrence + gated-RMS-norm decode.
gdn_decode_f32_norm: FP32 GDN decode fused with gated RMS norm.
gdn_decode_f32_norm_snap: super::gdn_decode_f32_norm + inline h-state snapshot.
gdn_decode_f32_strided: Strided FP32 GDN decode for concurrent sequence decode.
gdn_decode_f32_strided_norm: Strided FP32 GDN decode fused with gated RMS norm.
gdn_decode_f32_strided_norm_snap: super::gdn_decode_f32_strided_norm + inline h-state snapshot, for the batched-verify arm at batch_size = n sequences.
gdn_decode_wy2: WY-chunkwise 2-token GDN decode (2-pass algorithm).
gdn_decode_wy3: WY-chunkwise 3-token GDN decode (2-pass algorithm).
gdn_decode_wy4: WY-chunkwise 4-token GDN decode (2-pass algorithm).
gdn_decode_wyn: WY-Chunkwise Gated Delta Rule, pool-layout intermediates — K-generic launch shared by the K=17 DFlash verify (gated_delta_rule_wy17) and the chain-verify K∈{5..8} instantiations (gated_delta_rule_wy5..wy8, one templated source gated_delta_rule_wyn.cu). K is compile-time in the kernel; the caller selects it via the kernel handle. Computes K H·k dot products in 1 pass over H, applies WY algebraic correction over K tokens (K*(K-1)/2 inter-token k-dots), then applies K state updates in a second fused pass writing Hi_0..Hi_{K-2} + final H.
gdn_prefill: Gated delta rule prefill (multi-token, sequential SSM update within kernel).
gdn_prefill_fla: FLA multi-kernel chunked GDN prefill (ATLAS_GDN_FLA=1).
gdn_prefill_persistent: Persistent GDN prefill — h_state stays in shared memory for entire sequence.
gdn_prefill_persistent_batched: Batched persistent GDN prefill (uses gated_delta_rule_prefill_persistent_batched or gated_delta_rule_prefill_persistent_wy4_batched).
gdn_prefill_persistent_smem: Persistent GDN prefill with explicit shared memory size. Used for WY4-persistent variant which needs more shared memory.
gdn_prefill_persistent_smem_batched: Batched WY32 persistent GDN prefill (uses gated_delta_rule_prefill_wy64_batched).
gdn_prefill_regresident: Register-resident token-sequential prefill recurrence (warm-replay path).
gdn_prefill_split: Split-v_dim prefill: 2 CTAs per v-head, 64 threads each.
gdn_prefill_split4: 4-way split prefill: 4 CTAs per v-head, 32 threads each (128 total CTAs).
gdn_prefill_split4_batched: Batched split4 GDN prefill (uses gated_delta_rule_prefill_split4_batched).
gdn_verify_fused_conv_k2: STAGE 1 fused K=2 MTP-verify conv1d+L2norm: both draft positions in one launch, with the position-0 conv-state snapshot written inline (replaces the per-token conv1d_update_l2norm ×2 + intervening copy_d2d).
gdn_verify_fused_conv_kn
gdn_verify_fused_conv_kn_batched: Fused generic-K DFlash-verify conv1d+L2norm: ALL K draft positions in one launch, with every per-token conv-state rollback snapshot written inline to a strided intermediates array (replaces the per-token conv1d_update_l2norm ×K + copy_d2d ×K sequence — 34 serialized ops at K=17). conv_state is left holding the committed (post final-position) window, which the kernel also duplicates as snapshot K-1, so the caller issues NO copies.
gdn_verify_fused_conv_kn_f32: FP32-output twin of super::gdn_verify_fused_conv_kn: one launch for all K verify positions of conv1d+SiLU+L2norm, FP32 conv rows (what the sequential-decode-exact GDN chain reads), every per-token conv-state rollback snapshot written inline. output_stride is in FP32 elements.
gdn_verify_fused_norm_k2: STAGE 1 fused K=2 MTP-verify gated-RMS-norm: both draft positions in one launch (replaces the per-token gated_rms_norm ×2). The Z gate is read from the deinterleaved [Q|K|V|Z] buffer at z_offset per position.
gemv_sw_from: Kill-switch polarity for lossless SW GEMV. ON unless ATLAS_NO_GEMV_SW is exactly "1". =0 does not disable (same == "1" reading as ATLAS_NO_LM_HEAD_BATCH_GEMV).
glm_hc_expand: Expand one BF16 hidden state into the hc_mult FP32 highway streams. First layer only.
glm_hc_post: glm_hc_post: out[j] = post[j] * block_out + Σ_i comb[i][j] * residual[i].
glm_hc_pre: hc_pre: collapse the hc_mult FP32 streams to one BF16 sequence and emit this site’s post / comb mixing coefficients.
grouped_gemm_mla: Grouped GEMM for MLA: G independent [M,K_g]@[N_g,K_g]^T→[M,N_g] in one launch. Grid: (M*G, ceil(N_g/4), 1) Block: (256, 1, 1)
hc_expand: Broadcast a single hidden state into hc_mult identical streams: streams[t, i, d] = hidden[t, d]. One block per token.
hc_head: Final collapse before the LM head: a single learned sigmoid-weighted sum over the hc_mult streams. One block per token.
hc_head_lowrank: The model-level mixer (use_combine=False): the same collapse with no injection vector.
hc_head_mean: Final collapse before the LM head: an unweighted mean over the hc_mult streams.
hc_head_site: The model-level final collapse before the LM head.
hc_post: Expand the sublayer output back into hc_mult streams, mixing the saved residual streams through the doubly-stochastic comb. out may alias residual. One block per token.
hc_post_lowrank: Inject the block output back into every stream: out[t, s*H + d] = residual[t, s*H + d] + block_out[t, d] * inj[t, s].
hc_post_site: Inject the block output back into every stream. out may alias residual.
hc_pre: Collapse hc_mult streams to one (RMS-rescaled mix → sigmoid pre weighted sum) and emit post / comb (Sinkhorn) for the matching hc_post. One block per token.
hc_pre_lowrank: Collapse the hc_mult streams to one, and emit the per-stream injection weights the matching hc_post_lowrank needs.
hc_pre_site: Collapse the streams to one and emit whatever the matching hc_post needs — post+comb for Sinkhorn, the injection vector in post_out for low-rank.
int8_gemm_faith2_prefill: int8 W4A8 prefill GEMM: requant BF16 activations to int8 (per-32 F32 scale) then C = (A_i8 * W_i8) folded with per-32 A/W block scales via int8_gemm_faith2. The weight is already int8 (see requant_w_nvfp4_int8).
l2_norm: L2 normalization (in-place): data[i] = data[i] / sqrt(sum(data^2) + eps).
log_cutlass_nvfp4_route
log_gemm_shape: Roofline instrumentation: log each unique (kernel, M, N, K) GEMM shape once, gated by ATLAS_GEMM_SHAPE_LOG=1. Used to cross-reference nsys per-call durations → achieved TFLOPS/bandwidth vs GB10 peak.
mamba2_ssd_bmm: K2: CB[c][g][t][s] = C_t . B_s (raw, fp32).
mamba2_ssd_cumsum: K1: per-chunk dt (softplus+clamp) and inclusive cumsum of the log-decay.
mamba2_ssd_scan: K3: fused chunk_state + state_passing + chunk_scan (h0 stays in shared memory, so the per-chunk states tensor vLLM round-trips through DRAM never exists).
mamba2_ssm_prefill: Mamba-2 SSM prefill: sequential recurrence across seq_len tokens in a single kernel.
mamba2_ssm_prefill_persistent: Persistent Mamba-2 SSM prefill: H in shared memory, reduces global traffic. Same parameters and launch config as mamba2_ssm_prefill.
mix_hc: mix_hc — the row count of hc_fn and hc_base: (2 + hc_mult) * hc_mult.
mla_batched_gemv: Paged decode attention (FP8 KV cache, single/multi sequence).
mla_cache_assemble: MLA cache assembly: fuse [kv_latent|k_rope]→K and [kv_latent|zeros]→V into 1 kernel.
mla_cache_assemble_batched: Batched MLA cache assembly for N tokens: K=[latent|rope], V=[latent|zeros]. 1 kernel replaces N*4 D2D copies+memsets per layer.
mla_fused_prefill: Fused MLA prefill: Q_absorption + attention + V_extraction in one kernel. Grid: (num_heads, seq_len, 1) Block: (256, 1, 1)
mla_kv_assemble_batched: Batched K/V assembly from kv_expanded + k_rope for N tokens. 1 kernel replaces Nnkv3 D2D copies per layer.
mla_paged_decode_fp8: MLA Paged Decode — FP8 variant for DeepSeek-V4-Flash with FP8 KV cache.
mla_paged_decode_nvfp4: MLA Paged Decode — NVFP4 variant for DeepSeek-V4-Flash.
mla_prefill_attention_320: MLA absorbed prefill attention (HDIM=320, simple scalar kernel). Grid: (num_q_heads, ceil(seq_len/16), batch) Block: (256, 1, 1)
mla_q_final_assemble_batched: Assemble Q_final from Q_absorbed + Q_rope: [absorbed|rope] per head per token.
mla_q_rope_extract_batched: Batched Q rope extract: [N, nq, hd] → [N, nq, rope] at offset nope per head. 1 kernel replaces N*nq D2D copies per layer.
mla_q_rope_scatter: MLA Q_rope scatter: copy rope portion from q_full to strided q_absorbed_buf. 1 kernel replaces 32 D2D copies.
mla_q_rope_writeback: MLA Q_rope writeback: scatter RoPE’d rope portions to strided layout. 1 kernel replaces 32 D2D copies.
mla_q_rope_writeback_batched: Batched Q rope writeback: [N, nq, rope] → [N, nq, hd] at offset nope per head.
moe_batched_blend: Batched sigmoid blend: output += sigmoid(dot(normed, gate_weight)) * shared_out.
moe_bf16_grouped_gemm: BF16 grouped GEMM for sorted MoE prefill (FP8-dequant-on-load path).
moe_build_tile_worklist: Build the compacted (expert, m_tile, n_tile) work-list for the persistent grouped-GEMM grid. Single-block, thread-0 serial — mirrors the moe_sort_by_expert launch style (grid [1,1,1], block [256,1,1]).
moe_decode_atomic_c4_finalize
moe_decode_atomic_c4_silu_down_accum
moe_expert_gate_up_shared: Fused gate+up expert GEMV with shared expert as extra blockIdx.y slot.
moe_expert_gate_up_shared_batch2: Fused gate+up expert GEMV for K=2 tokens with shared expert.
moe_expert_gate_up_shared_batch3: Fused gate+up expert GEMV for K=3 tokens with shared expert.
moe_expert_gate_up_shared_batch2_t: NVFP4 fused gate+up GEMV (transposed). K=2 batch.
moe_expert_gate_up_shared_batch3_t: NVFP4 fused gate+up GEMV (transposed). K=3 batch.
moe_expert_gate_up_shared_bf16: Fused gate+up expert GEMV with shared expert for BF16 weights.
moe_expert_gate_up_shared_bf16_batch2: Fused gate+up expert GEMV with shared expert for BF16 weights — K=2 batch.
moe_expert_gate_up_shared_fp8: Fused gate+up expert GEMV with shared expert for FP8 weights.
moe_expert_gate_up_shared_fp8_batch2: FP8 fused gate+up GEMV for batch=2 (MTP K=2 verify). Grid: (ceil(N/8), 2*(top_k+1), 2) Block: (128, 1, 1)
moe_expert_gate_up_shared_fp8_batch3: FP8 fused gate+up GEMV for batch=3 (MTP K=3 verify). Grid: (ceil(N/8), 3*(top_k+1), 2) Block: (128, 1, 1)
moe_expert_gate_up_shared_fp8_batch2_t: FP8 fused gate+up GEMV (transposed). K=2 batch.
moe_expert_gate_up_shared_fp8_batch3_t: FP8 fused gate+up GEMV (transposed). K=3 batch.
moe_expert_gate_up_shared_fp8_t: FP8 fused gate+up GEMV (transposed weight). Single-token decode.
moe_expert_gate_up_shared_prefill: Fused gate+up expert GEMV for N-token prefill with shared expert.
moe_expert_gate_up_shared_t: NVFP4 fused gate+up GEMV (transposed weight). Single-token decode.
moe_expert_gemv: Batched MoE expert W4A16 GEMV: runs top_k expert GEMVs in one launch.
moe_expert_gemv_gate_up: Fused gate+up expert GEMV: both projections in one kernel launch.
moe_expert_gemv_gate_up_2x: Register-tiled fused gate+up expert GEMV: 2 output rows per thread.
moe_expert_gemv_silu_down: Fused SiLU+down expert GEMV: computes silu(gate)*up inline as activation.
moe_expert_gemv_silu_down_2x: Register-tiled fused SiLU+down expert GEMV: 2 output rows per thread.
moe_expert_silu_down_shared: Fused SiLU+down expert GEMV with shared expert as extra blockIdx.y slot.
moe_expert_silu_down_shared_batch2: Fused SiLU+down expert GEMV for K=2 tokens with shared expert.
moe_expert_silu_down_shared_batch3: Fused SiLU+down expert GEMV for K=3 tokens with shared expert.
moe_expert_silu_down_shared_batch2_t: NVFP4 fused SiLU+down GEMV (transposed). K=2 batch.
moe_expert_silu_down_shared_batch3_t: NVFP4 fused SiLU+down GEMV (transposed). K=3 batch.
moe_expert_silu_down_shared_bf16: Fused SiLU+down expert GEMV with shared expert for BF16 weights.
moe_expert_silu_down_shared_bf16_batch2: Fused SiLU+down expert GEMV with shared expert for BF16 weights — K=2 batch.
moe_expert_silu_down_shared_fp8: Fused SiLU+down expert GEMV with shared expert for FP8 weights.
moe_expert_silu_down_shared_fp8_batch2: FP8 fused SiLU+down GEMV for batch=2. Grid: (ceil(N/8), 2*(top_k+1), 1) Block: (128, 1, 1)
moe_expert_silu_down_shared_fp8_batch3: FP8 fused SiLU+down GEMV for batch=3. Grid: (ceil(N/8), 3*(top_k+1), 1) Block: (128, 1, 1)
moe_expert_silu_down_shared_fp8_batch2_t: FP8 fused SiLU+down GEMV (transposed). K=2 batch.
moe_expert_silu_down_shared_fp8_batch3_t: FP8 fused SiLU+down GEMV (transposed). K=3 batch.
moe_expert_silu_down_shared_fp8_t: FP8 fused SiLU+down GEMV (transposed weight). Single-token decode.
moe_expert_silu_down_shared_prefill: Fused SiLU+down expert GEMV for N-token prefill with shared expert.
moe_expert_silu_down_shared_t: NVFP4 fused SiLU+down GEMV (transposed weight). Single-token decode.
moe_fp8_grouped_gemm: FP8 grouped GEMM for sorted MoE prefill — grid-compaction over the COMPACTED work-list built by moe_build_tile_worklist. THE routed-expert FP8 prefill kernel.
moe_fp8_grouped_gemm_ptrtable_n128: FP8-A pointer-table grouped GEMM with transposed NVFP4 weights.
moe_gate_topk_fused: Fused gate GEMV + topK softmax for M=1 decode.
moe_grouped_down_cutlass: Single-launch CUTLASS grouped NVFP4 DOWN projection. a is the post-SiLU intermediate [total_expanded, inter] (already expert-contiguous — no gather). host is the down third of the load-time snapshot; expert_offsets is the device i32 [num_experts+1] prefix sum.
moe_grouped_gate_up_cutlass: Single-launch CUTLASS grouped NVFP4 fused gate_up GEMM (Phase-2).
moe_hash_route: GPU-side MoE hash routing (DeepSeek-V4 hash_moe layers).
moe_hash_route_batched: Batched GPU-side MoE hash routing (DeepSeek-V4 hash_moe layers, prefill).
moe_permute_tokens: Gather token rows into expert-sorted order: permuted[i] = hidden[sorted_token_ids[i]]. permuted is [total_expanded, hidden]. One block per output row, threads stride over hidden. Used by the FP4 grouped gate_up path (the CUTLASS escape-hatch needs contiguous per-expert rows; the FP8 fused kernel gathers internally so it doesn’t need this).
moe_silu_mul: Element-wise SiLU activation + multiply: output[i] = silu(gate[i]) * up[i].
moe_sort_by_expert
moe_topk_sigmoid
moe_topk_sigmoid_batched: Batched sigmoid + correction-bias top-K MoE routing.
moe_topk_softmax: GPU-side MoE top-K softmax.
moe_topk_softmax_batched: Batched top-K softmax: N tokens in parallel.
moe_topk_softmax_bias: GPU-side MoE top-K sigmoid routing (Nemotron-H).
moe_topk_softmax_bias_batched: Batched twin: one block per token.
moe_topk_sqrtsoftplus: GPU-side MoE top-K sqrtsoftplus routing (DeepSeek-V4).
moe_topk_sqrtsoftplus_batched: Batched sqrtsoftplus + correction-bias routing (DeepSeek-V4 prefill).
moe_transpose_u8_batched
moe_unpermute_reduce_indexed: Unpermute + weighted reduce with pre-built reverse map.
moe_w4a4_grouped_gemm_relu2: Grouped W4A4 expert UP GEMM with fused relu^2 (native FP4 tensor cores). A is the pre-quantized NVFP4 latent (packed E2M1 + per-16 E4M3 scales); B comes from the per-expert NVFP4 pointer tables unchanged. Grid: (ceil(n_out/128), max_m_tiles, num_experts) Block: (128, 1, 1)
moe_w4a16_fused_gate_up_k64_m128: K64 fused gate+up GEMM — M=128 variant (Block D #3 — Avarok pattern).
moe_w4a16_fused_gate_up_k64_n128: K64 fused gate+up GEMM — zero pipeline stall for K=h (2048 for 35B), 32 K-steps vs 64.
moe_w4a16_fused_gate_up_n128: Fused gate+up grouped GEMM — single launch for both projections.
moe_w4a16_grouped_gemm: MoE grouped GEMM: per-expert W4A16 matrix multiply.
moe_w4a16_grouped_gemm_ptrtable: Pointer-table grouped GEMM: one launch covers all experts.
moe_w4a16_grouped_gemm_ptrtable_k64_n128: K64 down GEMM: K_STEP_T=64 eliminates pipeline stall (compute=128 cycles > load ~100 cycles). Use when K=inter (512 for 35B) — 8 K-steps vs 16 with K32.
moe_w4a16_grouped_gemm_ptrtable_m256: moe_w4a16_grouped_gemm_ptrtable with M_TILE=256 (512-thread block, 16 warps). Caller must pass max_m_tiles computed against 256, not 64 — see the div_ceil(4) at the call site, mirroring the m128 variant’s div_ceil(2).
moe_w4a16_grouped_gemm_ptrtable_n128: Pointer-table grouped GEMM with N_TILE=128 (transposed or wide kernels).
moe_w8a8_grouped_gemm: W8A8 + FP32 epilogue grouped MoE GEMM (vLLM-equivalent).
moe_w8a8_grouped_gemm_pm4: W8A8 + FP32 epilogue grouped MoE GEMM — PM4 geometry over the COMPACTED work-list built by moe_build_tile_worklist (kernel moe_w8a8_grouped_gemm_pm4, same module/numerics as moe_w8a8_grouped_gemm: bit-identical output, measured).
moe_weighted_sum_blend: Fused SiLU+down expert GEMV, wide variant (16 outputs/block for small K).
moe_weighted_sum_blend_batch2: Fused weighted sum + sigmoid blend for K=2 tokens.
moe_weighted_sum_blend_batch3: Fused weighted sum + sigmoid blend for K=3 tokens.
moe_weighted_sum_blend_prefill: Fused weighted sum + sigmoid blend for N-token prefill.
moe_zero_expert_add: out[t, :] += zero_accum[t] * x[t, :] — the identity-expert blend.
native_q2_mmq_enabled: Sub-flag gating the native Q2_0 MMQ prefill path (ATLAS_GGUF_NATIVE_Q2_MMQ=1). Default off: keep the transient-dequant stopgap so the two can be A/B’d on GPU. (ATLAS_GGUF_NATIVE_Q2 still gates keep-packing overall — this only chooses how the kept-packed weight is consumed in PREFILL.)
nvfp4_mmq_gemm: NVFP4 W4A4 MMQ GEMM: C[m,n] (bf16, missing ×scale2) = A_fp4[m,k] x W_nvfp4[n,k].
nvfp4_mmq_gemm_tiled: NVFP4 W4A4 MMQ GEMM with an M-SIZED TILE.
nvfp4_mmq_quantize_act: Quantize bf16 activations [m, k] -> block_fp4_mmq (e2m1 + ue4m3 group-16, ±2 scale search) into out_y. One thread per 16-value group; ne0 padded to 256.
nvfp4_mmq_repack: Repack a checkpoint NVFP4 weight (packed E2M1 [n, k/2] low=even/high=odd + E4M3 [n, k/16] scales) into llama block_nvfp4 rows [n][k/64]. Raw bit shuffle — the e2m1 codes and e4m3 scale bytes are reused verbatim (scale2 folded downstream).
nvfp4_mmq_smem: Dynamic shared memory for a given M-tile, from the vendor’s layout: ids_dst[mmq_x] + y-tile[mmq_x * MMQ_TILE_Y_K(=36) ints, padded to 256] + x-tile[128 * MMQ_MMA_TILE_X_K_FP4(=76) ints], all 4-byte. Reproduces the previously-hardcoded 57856 at mmq_x=128, which is the check that this derivation matches the kernel’s actual layout.
nvfp4_mmq_weight_bytes: Bytes for the block_nvfp4 form of an [n, k] weight (k % 64 == 0).
nvfp4_scale_bf16: In-place ×scale2 for the down-projection MMQ output ([m, h] bf16).
nvfp4_silu_mul_quant: Fused SiLU-mul + block_fp4_mmq quantize for the down-MMQ path: reads RAW gate/up MMQ outputs, applies the scale2 folds + swiglu clamp + SiLU-mul, and quantizes straight into the down GEMM’s y-format — the intermediate bf16 activation tensor is never written (this round-trip is why the unfused down arm measured neutral).
nvfp4_silu_mul_scaled: SiLU(gate×gs)×(up×us) with the per-projection scale2 fold (swiglu ±10 clamp, mirrors moe_silu_mul). In-place safe (out may alias gate).
paged_decode_attn_bf16
paged_decode_attn_bf16k_turbo2v: Paged decode attention for Bf16K + Turbo2V asymmetric KV cache (6.4x V comp).
paged_decode_attn_bf16k_turbo3v: Paged decode attention for Bf16K + Turbo3V asymmetric KV cache.
paged_decode_attn_bf16k_turbo4v: Paged decode attention for Bf16K + Turbo4V asymmetric KV cache.
paged_decode_attn_fp8
paged_decode_attn_fp8k_turbo2v: Paged decode attention for Fp8K + Turbo2V asymmetric KV cache (6.4x V comp).
paged_decode_attn_fp8k_turbo3v: Paged decode attention for Fp8K + Turbo3V asymmetric KV cache.
paged_decode_attn_fp8k_turbo4v: Paged decode attention for Fp8K + Turbo4V asymmetric KV cache.
paged_decode_attn_nvfp4: Paged decode attention (NVFP4 KV cache, single/multi sequence).
paged_decode_attn_reduce_fp8: Reduce split-K partials into final BF16 output (FP8 variant).
paged_decode_attn_reduce_nvfp4: Reduce split-K partials into final BF16 output.
paged_decode_attn_splitk_fp8: Split-K paged decode attention (FP8 KV cache).
paged_decode_attn_splitk_nvfp4: Split-K paged decode attention (NVFP4 KV cache).
paged_decode_attn_turbo3k_turbo8v: Paged decode attention for Turbo3K + Turbo8V asymmetric KV cache.
paged_decode_attn_turbo4k_turbo3v: Paged decode attention for Turbo4K + Turbo3V asymmetric KV cache.
paged_decode_attn_turbo4k_turbo8v: Paged decode attention for Turbo4K + Turbo8V asymmetric KV cache.
per_token_group_quant_fp8: Per-token-per-128-K-group FP8 activation quantization. Output: A_fp8 [M, K] FP8 E4M3 + a_scale [M, K/128] FP32. Matches vLLM’s per_token_group_quant_fp8.
ple_add_highway: highway += ple_out, in FP32. The reference adds PLE’s output to the residual before that layer’s attention hyper-connection.
ple_conv: Depthwise causal conv, kernel k_size, dilation dilation, plus the SiLU and the residual add against the un-normalized gated value.
ple_gate: Gate the n-gram value by the highway, and emit both the gated value and its norm_conv’d twin.
predequant_nvfp4_to_fp8
prefill_attention: Flash Attention v2 prefill on contiguous Q/K/V.
prefill_attention_64: Contiguous prefill Flash Attention — BF16, BR=64 (256 threads).
prefill_attention_512_sink: DeepSeek-V4 full-attention (non-CSA) prefill with a per-head attention sink.
prefill_attention_fp8kv: Contiguous prefill Flash Attention — FP8 E4M3 K/V variant (BR=64).
prefill_attention_paged: Paged prefill Flash Attention — reads K/V from paged KV cache via block_table.
prefill_attention_paged_64: Paged prefill Flash Attention — BF16 KV cache, BR=64 (256 threads).
prefill_attention_paged_512: Paged prefill Flash Attention for HDIM=512 (Gemma-4 full-attention) — BF16 KV.
prefill_attention_paged_batched: Batched BF16-KV paged prefill attention (BR=32).
prefill_attention_paged_batched_64: Batched BF16-KV paged prefill attention (BR=64, 256-thread variant).
prefill_attention_paged_bf16k_turbo2v_64: Prefill paged attention — TurboQuant+ safer-asym Bf16K + Turbo2V (BR=64).
prefill_attention_paged_bf16k_turbo3v_64: Paged prefill (BR=64) for Bf16K + Turbo3V asymmetric KV cache.
prefill_attention_paged_bf16k_turbo4v_64: Prefill paged attention — TurboQuant+ safer-asym Bf16K + Turbo4V (BR=64).
prefill_attention_paged_dflash: DFlash γ-block paged Flash Attention — BF16 KV cache variant.
prefill_attention_paged_dflash_bf16_indirect: DFlash γ-block paged Flash Attention — BF16 KV cache, INDIRECT scalar args.
prefill_attention_paged_fp8: Paged prefill Flash Attention — FP8 KV cache variant.
prefill_attention_paged_fp8_64: Paged prefill Flash Attention — FP8 KV cache, BR=64 (256 threads).
prefill_attention_paged_fp8_batched: Batched FP8-KV paged prefill attention (BR=32).
prefill_attention_paged_fp8_batched_64: Batched FP8-KV paged prefill attention (BR=64, 256-thread variant).
prefill_attention_paged_fp8_dflash: DFlash γ-block paged Flash Attention — FP8 KV cache variant.
prefill_attention_paged_fp8k_turbo2v_64: Prefill paged attention — TurboQuant+ asym Fp8K + Turbo2V (BR=64, 6.4x V comp).
prefill_attention_paged_fp8k_turbo3v_64: Prefill paged attention — TurboQuant+ asym Fp8K + Turbo3V (BR=64).
prefill_attention_paged_fp8k_turbo4v_64: Prefill paged attention — TurboQuant+ asym Fp8K + Turbo4V (BR=64).
prefill_attention_paged_nvfp4: Paged prefill Flash Attention — NVFP4 KV cache variant.
prefill_attention_paged_nvfp4_64
prefill_attention_paged_nvfp4_batched: Batched NVFP4-KV paged prefill attention (BR=32).
prefill_attention_paged_nvfp4_batched_64: Batched NVFP4-KV paged prefill attention (BR=64, chunk_len >= 256).
prefill_attention_paged_turbo2_64
prefill_attention_paged_turbo3k_turbo8v_64: Prefill paged attention — TurboQuant+ asym Turbo3K + Turbo8V (BR=64).
prefill_attention_paged_turbo4k_turbo3v_64: Prefill paged attention — TurboQuant+ asym Turbo4K + Turbo3V (BR=64).
prefill_attention_paged_turbo4k_turbo8v_64: Prefill paged attention — TurboQuant+ asym Turbo4K + Turbo8V (BR=64).
prefill_attention_paged_turbo_64: Paged prefill Flash Attention — symmetric TurboQuant KV cache, BR=64. Shared launch wrapper for the turbo8 / turbo4 / turbo3 _64 kernel entries: identical ABI, the caller selects the dtype via kernel and passes that pool’s block stride + data-section offset.
prefill_batched_first_chunk_enabled: Whether chunk-zero streams may use the paged batched-prefill path.
prefill_varlen_enabled: VARLEN (ragged) batched prefill enabled? (--prefill-varlen-batch, legacy ATLAS_PREFILL_VARLEN=1; default OFF).
q2_0_gemv: Q2_0 GEMV (M=1 decode): C[1,N] = A[1,K] @ dequant(B), weights kept packed.
q2_0_gemv_batchm: Q2_0 batched GEMV (M=1..8 decode): C[M,N] = A[M,K] @ dequant(B).
q2_0_gemv_vec: Q2_0 GEMV (M=1 decode), CANDIDATE B: C[1,N] = A[1,K] @ dequant(B).
q2_0_gemv_vec_batchm: Q2_0 batched GEMV (M>=1 decode), CANDIDATE B: C[M,N] = A[M,K] @ dequant(B).
q2_0_mmq_gemm: Q2_0 MMQ GEMM: C[m,n] (bf16) = A_q8[m,k] x W_q2_0[n,k]. Fused bf16 store.
q2_0_mmq_gemm_packed: Q2_0 MMQ GEMM against a PackedQ2Weight (asserts group == 128, the only group the MMQ block layout supports — callers fall back to transient-dequant for group 64). Convenience over q2_0_mmq_gemm.
q2_0_weight_bytes: Bytes for the packed block_q2_0 form of an [n, k] weight (k % 128 == 0).
q4k_mmq_gemm: Q4_K MMQ GEMM: C[m,n] (bf16) = A_q8[m,k] x W_q4k[n,k]. Fused bf16 store.
q4k_weight_bytes: Bytes for the Q4_K-quantized form of an [nrows, n_per_row] weight (n_per_row % 256 == 0).
q8_1_scratch_bytes: q8_1_mmq activation scratch bytes for [m, k]; generous (kpad rounded to 256).
qsa_block_pool: Pool n_new freshly complete blocks starting at first_block: mean over ratio raw keys -> RMSNorm*(1+w) -> rope at block-start pos.
qsa_gather: Pack the selected tokens’ K/V rows into contiguous NHD scratch.
qsa_prefill_attn: Stage 2: per-row selected-set attention, overwriting the context rows.
qsa_qprep: One decode query: per head, RMSNorm*(1+w) + partial rope at pos -> FP32.
qsa_qprep_rows: Stage 2: per-row q prep for a contiguous selective row range.
qsa_score: scores[b] = sum_h relu(q_h . k_b) / sqrt(hd) over n_blocks blocks.
qsa_score_rows: Stage 2: per-row block scores, -inf beyond each row’s complete count.
qsa_score_rows_tc: Tensor-core qsa_score_rows (split-q). Geometry differs from the per-(row,block) kernel: one CTA covers 16 rows x 64 blocks with 8 warps, so the launch count drops from rows*blocks to ~1/1000th of that.
quant_gemm: Unified GEMM dispatch: select kernel based on weight quantization format.
quant_gemv: Unified GEMV dispatch: select kernel based on weight quantization format.
quantize_act_q8_1: Quantize bf16 activations [m, k] -> q8_1_mmq (DS4 layout) into out_q8.
quantize_bf16_to_fp8: Quantize a BF16 weight matrix [N, K] to FP8 E4M3 [N, K] with per-row f32 scales [N]. One CTA per row, 256 threads — parallel absmax reduction over K, then per-element saturating cast to E4M3.
quantize_bf16_to_nvfp4: Quantize a BF16 [M, K] matrix to NVFP4 (single-level, scale2=1.0): packed E2M1 [M, K/2] + per-group-16 E4M3 scales [M, K/16]. Prepares W4A4 prefill activations. Grid = M rows (one block/row), block 128 (threads stride groups).
quantize_weight_q4k: Quantize bf16 weights [nrows, n_per_row] -> GGML block_q4_K (at model load).
read_expert_ptrs_u64: One blocking D2H of a device [n] u64 pointer table. Load-time only.
read_expert_scales_f32: One blocking D2H of a device [n] f32 scale table. Load-time only.
requant_w_nvfp4_int8: Requant an NVFP4 weight (packed E2M1 + per-16 E4M3 block scales + per-tensor scale2) into an int8 weight + per-32 F32 block scale, for the int8 W4A8 prefill GEMM (int8_gemm_faith2). One-time conversion per weight at load (or lazily on first int8 prefill).
reshape_and_cache: Write K/V to paged FP8 cache using slot_mapping.
reshape_and_cache_bf16k_turbo2v: Write K/V to paged Bf16K + Turbo2V (TurboQuant+ safer-asym) cache (6.4x V comp).
reshape_and_cache_bf16k_turbo3v: Write K/V to paged Bf16K + Turbo3V (TurboQuant+ safer-asym) cache.
reshape_and_cache_bf16k_turbo4v: Write K/V to paged Bf16K + Turbo4V (TurboQuant+ safer-asym) cache.
reshape_and_cache_flash_v_only: V-only paged cache write — companion to the fused K-path so the K side of the cache stays exclusively owned by fused_k_norm_rope_cache_write_*. Use this when the fused K kernel is active to avoid the existing reshape_and_cache overwriting the correct K values with a double-rounded copy.
reshape_and_cache_fp8: k_cache/v_cache are the full pool base pointers. cache_stride is in elements (block_size * num_kv_heads * head_dim).
reshape_and_cache_fp8k_turbo2v: Write K/V to paged Fp8K + Turbo2V (TurboQuant+ asym) cache (6.4x V comp).
reshape_and_cache_fp8k_turbo3v: Write K/V to paged Fp8K + Turbo3V (TurboQuant+ asym) cache.
reshape_and_cache_fp8k_turbo4v: Write K/V to paged Fp8K + Turbo4V (TurboQuant+ asym) cache.
reshape_and_cache_nvfp4: Write K/V to paged NVFP4 cache (E2M1 data + per-group FP8 scales).
reshape_and_cache_turbo3k_turbo8v: Write K/V to paged Turbo3K + Turbo8V cache.
reshape_and_cache_turbo4k_turbo3v: Write K/V to paged Turbo4K + Turbo3V cache.
reshape_and_cache_turbo4k_turbo8v: Write K/V to paged Turbo4K + Turbo8V cache.
residual_add: BF16 residual add: residual[i] += src[i] (in-place).
residual_add_rms_norm: Fused residual add + RMS norm + residual save.
residual_add_rms_norm_gatef32: Dual-output fused residual add + RMS norm (ATLAS_FP32_ROUTING).
rms_norm
rms_norm_residual: Fused RMS norm + residual save: normed = rms_norm(input), residual = input.
rms_norm_short_row_eligible: Gate for rms_norm_warp_row: short even rows, many of them. Disable with ATLAS_RMS_NORM_WARP_ROW=0.
rms_norm_strided: RMS normalization: output = rms_norm(input) * weight.
rms_norm_warp_row: Warp-per-row RMS norm for SHORT rows — one warp per row instead of one block, so the grid shrinks 8x and the reduction needs no shared memory or barrier. Profitable exactly for the Qwen3 per-head q_norm/k_norm during prefill (num_rows = heads * seq, hidden_size = head_dim), where the block-per-row kernel measured ~43x above its bandwidth floor.
rope
rope_mrope_interleaved: MRoPE (interleaved multi-modal rotary) for Qwen3.6.
rope_mrope_interleaved_k_only: MRoPE for K only. Used when Q was already rotated by a fused Q prefill kernel.
rope_proportional: Proportional RoPE (Gemma-4 full-attention layers).
rope_strided: RoPE: apply rotary position embeddings to Q and K in-place.
rope_yarn: RoPE with precomputed YaRN inv_freq table (Mistral Small 4). The kernel reads frequencies from the table instead of computing from theta.
rope_yarn_scaled: YaRN RoPE for standard attention with explicit cosine/sine amplitude.
scaled_add: BF16 scaled accumulate: output[i] += scale * src[i].
set_prefill_varlen_from_cli: Publish the command line’s --prefill-varlen-batch decision. Returns the value IN FORCE, which differs from enabled when something already resolved the cell (then the command line did NOT take effect — the caller warns, mirroring gdn_flags::set_from_cli). Absent flag ⇒ never called ⇒ the documented ATLAS_PREFILL_VARLEN fallback stays reachable.
sigmoid_blend: Sigmoid-gated blend: output = output + sigmoid_gate * src.
sigmoid_blend_device: Sigmoid-gated blend reading gate scalar from device memory.
sigmoid_gate_mul: Element-wise sigmoid gate: output[i] = input[i] * sigmoid(gate[i]).
sigmoid_gate_mul_batched: Batched sigmoid gate multiply across multiple tokens.
sigmoid_gate_mul_head_broadcast: Per-head sigmoid gate multiply with broadcast over head_dim.
silu_mul: Fused SiLU activation: output = SiLU(gate) * up.
silu_mul_quant_fp8: Fused SiLU·mul + per-token-group(128) FP8-E4M3 quantization — replaces the silu_mul → per_token_group_quant_fp8 pair on the W8A8 prefill down-path without materializing the BF16 intermediate. Bit-identical to the pair (product rounds through BF16 before the group max; same reduction order, scale floor, and SATFINITE encode).
softplus_gate_mul_head_broadcast: Per-head softplus gate multiply with broadcast over head_dim.
ssd_scan_fits: Whether the SSD chunked scan physically fits for this state_size.
ssd_scan_smem: Dynamic shared memory mamba2_ssd_scan requests for a given SSM state size. Grows ~linearly in state_size, so a checkpoint with a large SSM state can exceed MAX_DYNAMIC_SMEM. Kept next to the launch so the two cannot drift.
ssm_h_state_f16_to_f32: One-shot FP16 -> FP32 widening of one layer’s SSM h-state (ATLAS_SSM_H_FP16). n is the FP32 ELEMENT count of the destination.
ssm_h_state_f32_to_f16: One-shot FP32 -> FP16 conversion of one layer’s SSM h-state (ATLAS_SSM_H_FP16). n is the FP32 ELEMENT count, derived from the pool’s byte size — never a duplicated shape literal.
transpose_block_scale: Transpose block scales: [N/128, K/128] → [K/128, N/128].
transpose_fp8: Transpose FP8 weight matrix on GPU: B[N,K] → B_t[K,N]. Grid: (ceil(N*K/256), 1, 1) Block: (256, 1, 1)
transpose_u8: Batched per-expert uint8 transpose for MoE down_proj relayout.
use_gemv_sw: SW kernel when the model lever is on and the handle resolved.
w4a4_gemm: W4A4 NVFP4 prefill GEMM (native FP4 tensor cores, sm_121a). Activation is pre-quantized NVFP4 (a_packed/a_scale, scale2=1.0); weight is the native NVFP4 QuantizedWeight. Output BF16 [M, N]. See kernels/…/w4a4_gemm.cu. Grid: (ceil(N/128), ceil(M/128), 1) Block: (256, 1, 1).
w4a4_gemm_mfast: W4A16 GEMM with N_TILE=128: same kernel signature, wider N tile.
w4a16_decode_gemv: Decode GEMV: software-pipelined single-warp when the lever and handle agree.
w4a16_gemm: W4A16 GEMM: C = A @ dequant(B).
w4a16_gemm_n128
w4a16_gemm_n128_ldb: W4A16 GEMM with N_TILE=128: same kernel signature, wider N tile.
w4a16_gemm_n128_m128: W4A16 GEMM: C = A @ B with 2-M-chunk CTA (M_TILE2=128).
w4a16_gemm_n128_m128_bf16: 8-arg launcher for w4a16_gemm_t_m128_bf16 (v1) ONLY. The _v2 sibling’s compiled signature has a 9th ldb param — launching it through this helper makes cuLaunchKernel read one-past-the-end of the param array (observed as CUDA_ERROR_INVALID_VALUE or a host SIGSEGV depending on the neighboring heap word). Launch v2 via w4a16_gemm_n128_m128_bf16_ldb (ldb = N when the transposed twin is unpadded).
w4a16_gemm_n128_m128_bf16_ldb: W4A16 GEMM — LOSSLESS BF16 prefill variant of w4a16_gemm_n128_m128.
w4a16_gemm_n128_m128_v2: W4A16 GEMM v2: shadow of w4a16_gemm_n128_m128 (minimax, step3p7, and — since the 27B port — qwen3.6-27b).
w4a16_gemm_n128_m128_v3: W4A16 GEMM v3: MiniMax-only shadow with K_STEP=64 (was 32 in v2). Halves K-iteration count; doubles per-iter MMA count. 1 CTA/SM (was 3 for v2) due to larger SMEM footprint.
w4a16_gemv: W4A16 GEMV (M=1): C = A @ dequant(B) for single-row activations.
w4a16_gemv_batch2: W4A16 double-GEMV (M=2): reads weights once, computes 2 outputs.
w4a16_gemv_batch3: W4A16 triple-GEMV (M=3): reads weights once, computes 3 outputs.
w4a16_gemv_batchm: W4A16 batched GEMV (M<=MAX_M) — the NVFP4 sibling of w8a16_gemv_batch4/16.
w4a16_gemv_dual: W4A16 dual GEMV: two projections sharing the same BF16 input, one launch.
w4a16_gemv_dual_batch2: Dual-projection GEMV for 2 tokens (K+V or any 2 weight matrices).
w4a16_gemv_dual_batch3: Dual-projection GEMV for 3 tokens (K+V or any 2 weight matrices).
w4a16_gemv_dual_sw: Single-warp-per-output variant of w4a16_gemv_dual (8 outputs/block → N/8 grid). Bit-identical output (see w4a16_gemv_fused.cu). Default ON via ModelLevers::gemv_sw; kill with ATLAS_NO_GEMV_SW=1.
w4a16_gemv_grid_x
w4a16_gemv_qg: W4A16 GEMV with inline Q/Gate deinterleave on output write.
w4a16_gemv_qg_batch2: Q+Gate GEMV for 2 tokens with inline deinterleave.
w4a16_gemv_qg_batch3: W4A16 GEMV batch3 with inline Q/Gate deinterleave.
w4a16_gemv_qkvz: W4A16 GEMV with inline QKVZ deinterleave on output write.
w4a16_gemv_silu_input: W4A16 GEMV with fused SiLU input: silu(gate)*up as activation, GEMV with down weights.
w4a16_gemv_silu_input_sw: Single-warp-per-output variant of w4a16_gemv_silu_input (N/8 grid). Bit-identical. Default ON via ModelLevers::gemv_sw.
w4a16_gemv_sw: Single-warp-per-output W4A16 GEMV (M=1). Grid: (ceil(N/8), 1, 1).
w4a16_gemv_sw_grid_x
w4a16_gemv_sw_raw: Same launch as w4a16_gemv_sw for callers that hold the NVFP4 operand triple as loose pointers rather than a QuantizedWeight (GLM-5.3’s Nvfp4Proj).
w8a16_gemm: W8A16 GEMM (M>1): C[M,N] = A[M,K] @ dequant(B[N,K]) for prefill.
w8a16_gemm_n128_m128: W8A16 transposed M128 GEMM (kernel w8a16_gemm_t_m128): FP8 E4M3 analog of w4a16_gemm_n128_m128_v2. 128×128 (M×N) tile, two 64-row chunks, 8 warps, parallel-chunk m16n8k16.bf16.bf16 MMA + two-level FP32 block-scale fold. Same transposed contract as w8a16_gemm_t (B_t[K,N] + block_scale_t[K/128, N/128]); reuses the transpose_fp8 / transpose_block_scale output as-is. Grid: (ceil(N/128), ceil(M/128), 1) Block: (256, 1, 1)
w8a16_gemm_pipelined: W8A16 GEMM pipelined (M>1): bit-identical (cosine=1.0) faster rewrite of w8a16_gemm — same args, same numerics, ~4.6× faster on GB10/sm_121.
w8a16_gemm_t: W8A16 Transposed GEMM: C[M,N] = A[M,K] @ dequant(B_t[K,N]) with coalesced reads.
w8a16_gemm_t_pipelined: Pipelined transposed W8A16 GEMM (kernel w8a16_gemm_t_pipelined): same transposed args as w8a16_gemm_t, ~4.2x via smem-LUT + K_STEP32 + K-contiguous smem_B + 128x32 occupancy tile. Grid: (ceil(N/32), ceil(M/128), 1) Block: (256, 1, 1)
w8a16_gemv: W8A16 GEMV (M=1): C = A @ dequant_lut(B_fp8) * row_scale for FP8 E4M3 weights.
w8a16_gemv_batch2: Block-scaled FP8 dual-GEMV (batch=2). input is [2, K] BF16, output is [2, N] BF16; weight/block_scale are the raw w8a16_gemv pointers. Grid: (ceil(N/4), 1, 1) Block: (256, 1, 1)
w8a16_gemv_batch4: Block-scaled FP8 batched GEMV (M<=4). input is [M, K] BF16, output is [M, N] BF16; weight/block_scale are the raw w8a16_gemv pointers (2D block-scaled FP8). One pass over the FP8 weight serves all M rows — the M=4 sibling of w8a16_gemv, replacing w8a16_gemm_pipelined for n<=4 batched decode (which pads M to a 128-row MMA tile). Bit-identical per-row to w8a16_gemv. Grid: (ceil(N/4), 1, 1) Block: (256, 1, 1)
w8a16_gemv_dual: W8A16 (FP8 E4M3) dual GEMV: two projections sharing the same BF16 input, one launch. blockIdx.z selects projection 0 (gate) vs 1 (up). Both N must be equal. Mirrors w4a16_gemv_dual but takes RAW DevicePtrs (FP8 weights are fp8w.weight / fp8w.row_scale, no QuantizedWeight wrapper, no scale2 f32).
w8a16_gemv_silu_input: W8A16 (FP8 E4M3) GEMV with fused SiLU input: silu(gate)*up as activation, GEMV with FP8 down weights. Reads gate_out[K] and up_out[K] BF16, computes silu(gate)*up per element inline, then multiplies by dequanted FP8 down weights. Eliminates the separate silu_mul kernel + down GEMV. Mirrors w4a16_gemv_silu_input but with RAW DevicePtrs (no scale2 f32).
wide_prefill_kernel: The HDIM>256 prefill kernel: its module/entry name and the BR its grid must be built for. ONE reader, because the name is chosen in qwen3_attention::init and the grid here, and a mismatch is silent.
widen_block_scale_f32: Widen an FP8 block-scale tensor to FP32 on the GPU.

Module ops

Module ops Copy item path

Re-exports§

Modules§

Structs§

Enums§

Constants§

Functions§

Module ops