metrale_model_arch/dflash_head/
kernels.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `DflashKernels`, the kernel handles the DFlash drafter launches.
4//!
5//! Owner: model-arch (DFlash drafter).
6//! Invariants: none beyond the types.
7
8use super::*;
9
10/// 2026-09-25: Kernel handles for the DFlash drafter, all resolved once in
11/// `BlockDiffusionDraftHead::from_weights`. A handle resolved there with
12/// `try_kernel` is `KernelHandle(0)` when the target's build lacks the kernel.
13pub struct DflashKernels {
14    pub rms_norm: KernelHandle,
15    pub residual_rms_norm: KernelHandle,
16    pub dense_gemv: KernelHandle,
17    pub dense_gemm: KernelHandle,
18    /// 2026-09-25: NVFP4 GEMM for the logits when the target's lm_head is NVFP4
19    /// (`lm_head_nvfp4` is `Some`), whose packed bytes the BF16 GEMM cannot
20    /// read. `.0 == 0` when the target has no `w4a16_gemm` kernel.
21    pub w4a16_gemm: KernelHandle,
22    pub dense_gemm_pipelined: KernelHandle,
23    pub rope_qwen3: KernelHandle,
24    pub reshape_cache_fp8: KernelHandle,
25    /// 2026-09-25: BF16 paged KV write: the ctx rows in `precompute_ctx_kv` and
26    /// the block rows in `forward_block_layer_pre_attn`.
27    pub reshape_cache_bf16: KernelHandle,
28    pub prefill_attn_dflash_fp8: KernelHandle,
29    pub prefill_attn_dflash_bf16: KernelHandle,
30    /// 2026-09-25: `attn_prefill_paged_indirect`: BF16 paged attention that reads
31    /// `kv_len`, `q_offset` and `q_rope_pos` from device memory at kernel entry.
32    /// The drafter's paged attention (`forward_block_layer_attention`) launches it
33    /// non-causally.
34    pub prefill_attn_dflash_bf16_indirect: KernelHandle,
35    pub silu_mul: KernelHandle,
36    pub residual_add: KernelHandle,
37    pub argmax: KernelHandle,
38    pub batched_embed: KernelHandle,
39    /// 2026-09-25: Builds `count` i64 cache slot indices on the device from a
40    /// block table: the ctx rows' slots in propose and the batched precompute,
41    /// and the block rows' slots in `forward_block`.
42    pub fill_slots: KernelHandle,
43    /// 2026-09-25: Non-paged prefill attention, built for head_dim 128 only
44    /// (`attn_prefill_h128`): the non-paged layer body and the
45    /// `dflash_contig_attn` path use it.
46    pub prefill_attn: KernelHandle,
47    /// 2026-09-25: BF16 → FP8 E4M3 weight quantization with one f32 scale per
48    /// row, run only in `from_weights` to build the FP8 mirrors of the seven
49    /// dense GEMM weights (q/k/v/o/gate/up/down) and, when the checkpoint has
50    /// no FP8 lm_head, the lm_head.
51    pub quantize_bf16_to_fp8: KernelHandle,
52    /// 2026-09-25: Row-scaled BF16 × FP8 → BF16 GEMM over an `Fp8DenseWeight`,
53    /// from the `w4a16` module (`from_weights` picks the variant). With
54    /// `DflashQuantization::Fp8Weights` the layer GEMMs use it when neither rt2
55    /// GEMV applies, and the lm_head uses it above 16 rows.
56    pub fp8_gemm_n128_row_scaled: KernelHandle,
57    pub dense_gemv_fp8w: KernelHandle,
58    /// 2026-09-25: `fp8_gemm_t_row_scaled_m16`, a single-warp row-scaled FP8 GEMM
59    /// that computes at most 16 rows. The FP8 lm_head uses it at 16 rows or fewer
60    /// when neither rt2 GEMV applies.
61    pub fp8_gemm_n128_row_scaled_m16: KernelHandle,
62    /// 2026-09-25: `fp8_gemv_rowscale_batch8_rt2`, a register-tiled row-scaled FP8
63    /// GEMV for at most 8 rows. The FP8 layer GEMMs and the FP8 lm_head try it
64    /// first, then `fp8_gemv_rt2_16`, then a tile GEMM. `.0 == 0` on targets
65    /// without the `fp8_gemv_rt` module.
66    /// Kill-switch: METRALE_NO_DFLASH_FP8_RT=1. provenance-id:
67    /// 526f6e616c6420522e205374657369616b
68    pub fp8_gemv_rt2: KernelHandle,
69    /// 2026-09-25: `fp8_gemv_rowscale_batch16_rt2`, the 16-row sibling of
70    /// `fp8_gemv_rt2`, for at most 16 rows. `.0 == 0` on targets without it.
71    /// Both rt2 GEMVs also need K to be a multiple of 16, and
72    /// `METRALE_NO_DFLASH_FP8_RT=1` disables both.
73    /// provenance-id: 526f6e616c6420522e205374657369616b
74    pub fp8_gemv_rt2_16: KernelHandle,
75    /// 2026-09-25: DFlash2 two-tap grouped dynamic conv (`kernels/gb10/common/dflash2.cu`).
76    /// `.0 == 0` on targets without the module, and `dflash2_active` is then false.
77    pub dflash2_conv2: KernelHandle,
78    /// 2026-09-25: DFlash2 per-row top-16 over the drafter logits; overwrites each
79    /// selected logit.
80    pub dflash2_topk16: KernelHandle,
81    /// 2026-09-25: DFlash2 candidate-selector chain walk, one launch, one block per
82    /// sequence.
83    pub dflash2_selector_walk: KernelHandle,
84}
metrale_model_arch/dflash_head/kernels.rs

metrale_model_arch/dflash_head/
kernels.rs
