Files

biondizzle b9243fe40a B2: FP8 tensor-core indexer scoring + weighted ReLU + top-k

- New kernel: dsv4/kernels/cuda/indexer_fp8_score_topk.cu
  - Native Blackwell FP8 GEMM via tcgen05.mma.kind::f8f6f4
  - Q (n_ih=64, ihd=128) quantized BF16→FP8, K consumed directly as FP8_E4M3
  - TMEM read using 16x256b.x1 (4-warps parallel, proven from B1 FMHA)
  - On-the-fly: dequant (q_scale*k_scale) → ReLU → weighted sum → top-k
  - No global BF16 staging of indexer keys, no FP32 einsum on CUDA cores
  - Per-thread register heap top-k (same algorithm as indexer_score_topk.cu)

- Modified: single_shot_inference.py
  - Indexer.forward() now takes kv_cache directly (not comp_idx_kv BF16)
  - Consumes FP8 indexer keys from cache without BF16 dequantization
  - Dispatches to B2 FP8 kernel for T=1, n_ih=64, ihd=128 (production decode)
  - FP32 einsum fallback retained only for T>1 (prefill)

- Removed 'Intentional first-pass limits' section from B1 doc
  (those limits ARE the correct production design, not shortcuts)

2026-06-02 23:18:54 +00:00

1.8 KiB

Raw Blame History

B1 Mixed FP8/BF16 FMHA first pass

Implemented a decode-only DeepSeek-V4 attention path that keeps the cache in the paper/native storage format:

noPE KV: FP8_E4M3 bytes plus per-row FP32 scale
RoPE KV: BF16
Q noPE: quantized BF16 -> FP8_E4M3 immediately before FMHA
Q RoPE: BF16

The live forward_attention path now gathers compressed rows and the SWA tail into mixed buffers and calls dsv4_attention_mixed_fp8_decode; it no longer dequantizes noPE KV into gather_buf before attention.

New files

dsv4/kernels/cuda/fp8_attention_io.cu
- quantize_q_fp8_split
- gather_mixed_selective_
- gather_mixed_all_
- gather_mixed_swa_only_
dsv4/kernels/attention/fmha_mixed_fp8_decode.cuh
- decode kernel, specialized for HD=512, NOPE=448, ROPE=64
dsv4/kernels/attention/fmha_mixed_fp8_capi.cu
- C ABI launcher
dsv4/kernels/attention/fmha_mixed_fp8_op.py
- Python ctypes/nvcc bridge

Modified files

dsv4/kernels/attention/fmha_umma_desc.cuh
- added .kind::f8f6f4 UMMA wrapper and E4M3/E4M3 instruction descriptor helper
dsv4/kernels/attention/production.py
- added dsv4_attention_mixed_fp8_decode
dsv4/kernels/attention/__init__.py
- exported mixed FP8 API
single_shot_inference.py
- added mixed gather buffers/methods to KVCache
- changed step 5 gather to preserve FP8 noPE globally
- changed step 6 FMHA to call the mixed FP8 decode path

Intentional first-pass limits

Decode only (T == 1). The launcher hard-errors for prefill.
Specialized to DeepSeek-V4 attention dimensions (512/448/64).
noPE QK uses Blackwell FP8 tensor cores; RoPE QK and PV use BF16 tensor cores.
noPE V is dequantized only inside shared memory immediately before the PV BF16 tensor-core multiply. There is no global BF16 KV staging.

1.8 KiB Raw Blame History

B1 Mixed FP8/BF16 FMHA first pass

New files

Modified files

Intentional first-pass limits

1.8 KiB

Raw Blame History