Struct DflashKernels 

pub struct DflashKernels {

    pub rms_norm: KernelHandle,
    pub residual_rms_norm: KernelHandle,
    pub dense_gemv: KernelHandle,
    pub dense_gemm: KernelHandle,
    pub w4a16_gemm: KernelHandle,
    pub dense_gemm_pipelined: KernelHandle,
    pub rope_qwen3: KernelHandle,
    pub reshape_cache_fp8: KernelHandle,
    pub reshape_cache_bf16: KernelHandle,
    pub prefill_attn_dflash_fp8: KernelHandle,
    pub prefill_attn_dflash_bf16: KernelHandle,
    pub prefill_attn_dflash_bf16_indirect: KernelHandle,
    pub silu_mul: KernelHandle,
    pub residual_add: KernelHandle,
    pub argmax: KernelHandle,
    pub batched_embed: KernelHandle,
    pub fill_slots: KernelHandle,
    pub prefill_attn: KernelHandle,
    pub quantize_bf16_to_fp8: KernelHandle,
    pub fp8_gemm_n128_row_scaled: KernelHandle,
    pub dense_gemv_fp8w: KernelHandle,
    pub fp8_gemm_n128_row_scaled_m16: KernelHandle,
    pub fp8_gemv_rt2: KernelHandle,
    pub fp8_gemv_rt2_16: KernelHandle,
    pub dflash2_conv2: KernelHandle,
    pub dflash2_topk16: KernelHandle,
    pub dflash2_selector_walk: KernelHandle,
}

Kernel handles for the DFlash γ-block forward chain. All resolved once at BlockDiffusionDraftHead::from_weights against the active GPU backend (which compiles target-specific PTX at startup); subsequent propose() calls just KernelLaunch::new(...).launch(stream).

Fields

rms_norm: KernelHandle

residual_rms_norm: KernelHandle

dense_gemv: KernelHandle

dense_gemm: KernelHandle

w4a16_gemm: KernelHandle

NVFP4 GEMM for the final logits when the shared lm_head is NVFP4 (e.g. Holo): a BF16 dense_gemm on NVFP4-packed bytes reads garbage (and ~4× OOB → CUDA-700). .0 == 0 when the target lm_head is BF16.

dense_gemm_pipelined: KernelHandle

rope_qwen3: KernelHandle

reshape_cache_fp8: KernelHandle

reshape_cache_bf16: KernelHandle

BF16 KV cache writeback. Used by Phase 2 precompute_ctx_kv and the per-layer γ-block reshape_and_cache call to populate the drafter’s BF16 paged cache before each prefill_attention_paged_dflash.

prefill_attn_dflash_fp8: KernelHandle

prefill_attn_dflash_bf16: KernelHandle

BF16 paged-attention dispatcher for the DFlash γ-block. Calls inferspark_prefill_paged with causal_mask_enabled=0, reading BF16 K/V from the per-layer paged cache pool. Phase 2 (Option B) drafter attention runs through this kernel; the FP8 variant above is retained for a future quality-validated FP8 KV path. See ops::prefill_attention_paged_dflash.

prefill_attn_dflash_bf16_indirect: KernelHandle

Phase 5 (CUDA graph) variant of prefill_attn_dflash_bf16 that reads kv_len and q_offset from device pointers instead of taking them as kernel scalar args. Used by the graph-captured forward_block path so a single graph instance can be replayed across steps with different dynamic values written to the indirect-args buffer pre-launch. Resolves to kernel inferspark_prefill_paged_indirect.

silu_mul: KernelHandle

residual_add: KernelHandle

argmax: KernelHandle

batched_embed: KernelHandle

fill_slots: KernelHandle

Phase 2 Option B: builds [count] i32 slot indices on-device from a host-provided block_table. Used by propose.rs to populate the slot_mapping passed to reshape_and_cache and precompute_ctx_kv.

prefill_attn: KernelHandle

Non-paged prefill attention (used for the γ-block self-attention when there’s no persistent K/V cache to walk).

quantize_bf16_to_fp8: KernelHandle

Phase G — BF16 → FP8 E4M3 per-row weight quantization. Used at model load time to convert the seven dense-GEMM drafter weights (q/k/v/o/gate/up/down) when ATLAS_DFLASH_DRAFTER_FP8=1. Never on the hot path.

fp8_gemm_n128_row_scaled: KernelHandle

Phase G — Row-scaled BF16 × FP8 → BF16 GEMM. Consumes the Fp8DenseWeight (FP8 weight + per-row f32 scale) produced at load time by quantize_bf16_to_fp8. Wraps kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu fp8_gemm_t_row_scaled. Replaces dense_gemm_bf16 on the seven dense-GEMM call sites in forward_block_layer_pre_attn / _post_attn when self.quant == DflashQuantization::Fp8Weights.

dense_gemv_fp8w: KernelHandle

Phase G — Row-scaled BF16 × FP8 → BF16 GEMV (M=1) for the lm_head fall-back. At γ=16 vs vocab=248320 the row-scaled GEMM wastes 75% of its M_TILE; the GEMV in a γ-loop is faster.

fp8_gemm_n128_row_scaled_m16: KernelHandle

Phase G — Small-M (M≤16) row-scaled FP8 GEMM. Drop-in replacement for fp8_gemm_n128_row_scaled when M=γ=16. Single warp per CTA, no wasted M_TILE rows. Used by the lm_head GEMM.

fp8_gemv_rt2: KernelHandle

Register-tiled batched row-scaled FP8 GEMV (M<=8, T=2 outputs per thread) — the FP8 twin of w4a16_gemv_batch8_rt2. Preferred over BOTH tile GEMMs above at M<=8 (they pad 87%/50% of their M-tile; ~100 GB/s measured vs 180+ for the rt family, nsys 2026-08-19). .0 == 0 on targets without the fp8_gemv_rt module → tile path. Kill-switch: ATLAS_NO_DFLASH_FP8_RT=1. provenance-id: 526f6e616c6420522e205374657369616b

fp8_gemv_rt2_16: KernelHandle

MAX_M=16 sibling of fp8_gemv_rt2 for the γ>8 propose window (2026-08-29: STEP_TIMING measured propose 18.2ms rt2 vs 38.0ms tile fallback at flag 9 — the entire γ>8 step tax). .0 == 0 on stale kernel builds → tile path, exactly as before. provenance-id: 526f6e616c6420522e205374657369616b

dflash2_conv2: KernelHandle

DFlash2 two-tap grouped dynamic conv (kernels/gb10/common/dflash2.cu). .0 == 0 on targets without the module (DFlash2 then refuses to arm).

dflash2_topk16: KernelHandle

DFlash2 per-row destructive top-16 over drafter logits.

dflash2_selector_walk: KernelHandle

DFlash2 candidate-selector chain walk (single launch, whole block).