DeepSeek-V4: FP4 CSA indexer (gfx950)#1581
Open
junhaha666 wants to merge 8 commits into
Open
Conversation
…ndexer Replace the ATOM_V4_INDEXER_FP4 env var with a first-class CLI flag. The flag flows EngineArgs -> engine kwargs -> Config (enable_deepseek_v4_fp4_indexer), and the attention metadata builder reads it instead of the environment. Since the FP4 mqa-logits / scatter kernels are gfx950 (MI355X / CDNA4) only, requesting the flag on any other arch logs a warning and falls back to the FP8 indexer rather than failing. The flag is also added to the config hash factors so toggling it picks a distinct compile/cache key. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…schedule compute_varctx_schedule is now a single Triton kernel, so build the FP4 decode indexer schedule straight into the fixed-address CUDAGraph buffer (cta_info_out=) instead of ~50 tiny torch launches + an intermediate alloc and .copy_. ~300 us -> ~40 us per decode step, which was the overhead making the FP4 indexer slower e2e than FP8 despite the faster mqa kernel. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # atom/config.py # atom/model_ops/attentions/deepseek_v4_attn.py # atom/models/deepseek_v4.py
The merge of origin/main brought in main's refactor that split the indexer's inline forward_batched into forward_pre (Q proj + RoPE + quant + weights, returns (q_fp8, weights)) + score_topk_from (the eager paged indexer_score_topk dispatch). main applied this to the FP8 path correctly, but the FP4 block inside forward_pre kept the pre-refactor ending `return torch.ops.aiter.indexer_score_topk(...)` (the final topk result), so both callers — forward_batched (`q_fp8, weights = forward_pre(...)`) and _attn_pre (PIECEWISE, `idx_q_fp8, idx_weights = forward_pre(...)`) — hit "ValueError: too many values to unpack (expected 2)" on the FP4 path. Return (q_fp4, weights) like the FP8 path; score_topk_from -> indexer_score_topk already routes FP4 via its uint8 kv_cache check and reads the stashed _q_scale_fp4. q_fp4 (uint8) rides the q_fp8 arg. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The FP4 indexer prefill allocated its seq-local logits at the model-max width (`_max_model_len_idx` = max_model_len//4, e.g. 262144), so a 16k prefill needed a [16384, 262144] fp32 buffer (~17 GB) even though each row only writes/reads [0, visible_end). Right-size the width to the batch's ACTUAL max committed index length (max over seqs of n_committed_csa), derived on the CPU from the already-host-side n_committed_csa_per_seq_cpu (no new device sync) — a 16k prefill now needs ~[16384, 4096] (~268 MB). Also factor the FP8 prefill's Q-row budget-chunking (introduced for #1376) into a shared `Indexer._prefill_chunked_topk` helper and route both the FP8 (GLOBAL-output) and FP4 (seq-local-output) prefill paths through it: - FP8: unchanged behavior, just deduplicated. - FP4: with the right-sized width, the [total_tokens, W] fp32 buffer stays within ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB in the common case, so the loop runs once and reuses the schedule precomputed ONCE (outside the fwd) by the metadata builder — no compute_prefill_schedule in the hot path. Chunking only engages for genuinely huge actual contexts, and only then does each chunk rebuild its schedule (cta_info encodes absolute row ids). Decode keeps the fixed `_max_model_len_idx` width (CUDA graph needs a static shape); prefill is eager so a per-fwd width is fine. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
The HCA compress paged-offset build assumed one HCA compressed entry per physical block (k2_hca==1, the old block_size=128). At V4 block_size=256, k2_hca = block_size // hca_ratio = 256 // 128 = 2, so each physical block packs TWO HCA entries — the compressor writes entry `ci` into its cache view [num_blocks, k2_hca, head_dim] at (block_tables[bid, ci // k2_hca], ci % k2_hca), i.e. unified row swa_pages + phys * k2_hca + slot. Both index builders emitted `swa_pages + block_tables[bid, ci]` (k2==1), so at k2_hca==2 half the HCA entries indexed the block table out past the seq's real blocks (reading unrelated/padding physical blocks) and dropped the intra-block slot — silently corrupting the HCA (128:1) layers' KV reads. It survives most end-to-end recall (CSA/SWA carry it) so it does not show up in output canaries, but it is a real correctness bug in the 20 HCA layers. Fix both paths with the k2-aware formula (matches the correct CSA csa_translate_pack `//capacity, %capacity, phys*capacity+slot` pattern): - prefill: `_v4_paged_prefill_indices_kernel` gets a `K2_HCA` constexpr; wrapper + numpy reference take `k2_hca`. - decode: extract the numpy offset math into a shared `hca_compress_paged_offsets` helper (in paged_decode_indices), used by `_attach_v4_paged_decode_meta`. - caller passes `k2_hca=self.k2_hca`. Regression tests (fail pre-fix, pass post-fix) at k2_hca=2: - test_prefill_indices_paged: kernel + reference + independent oracle. - test_decode_indices_paged: hca_compress_paged_offsets (+ k2==1 reduces to the legacy formula). No perf impact: index building is a per-fwd setup pass (not the attention hot path); the added arithmetic is constexpr div/mod by a power of 2 and the number of block-table loads is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an FP4 path for the DeepSeek-V4 CSA indexer (gfx950 / MI355X) as an
alternative to the existing FP8 indexer. When enabled, the indexer KV is stored
as packed FP4 (E2M1) + per-group(32) e8m0 scale and scored by the FlyDSL
pa_mqa_logits_fp4{,_prefill}kernels; Q is FP4-quantized inline. The FP8 pathis unchanged and remains the default.
Opt-in via
--enable-deepseek-v4-fp4-indexer(gfx950 only — other arches log awarning and fall back to FP8). Runtime scoring auto-selects by cache dtype
(uint8 FP4 cache → FP4 path; fp8 cache → FP8 path).
What's included
layout); Q FP4 quant fused into
rope_rotate_activation; Compressor writes viafused_compress_attn(quant_mode="fp4")._score_topk_{prefill,decode}_fp4over the paged FP4 cache viablock_tables(no cp_gather/deepgemm). Decode is CUDAGraph-safe — thepersistent-grid schedule (
cta_info) is precomputed into a fixed buffer via afused single-kernel schedule; prefill stays eager.
--enable-deepseek-v4-fp4-indexer+ config field + hash factor;gfx950 check with warn-and-fallback.
the batch's ACTUAL max committed index length (CPU-derived, no new sync)
instead of the model max — a 16k prefill drops from ~17 GB to ~268 MB. The
FP8 prefill's Q-row budget-chunking (GLM-5.2-FP8 OOM in sparse_attn_indexer (fp8_mqa_logits) under high concurrency + ~115k context #1376) is factored into a shared
Indexer._prefill_chunked_topkhelper used by both paths; with the right-sizedwidth FP4 runs single-chunk and reuses the once-precomputed schedule (no
compute_prefill_schedulein the hot path), chunking only for huge contexts.cache_scaleper-layer aliasing fix from the mainmerge.
Validation
fused_compress_attn(quant_mode="fp4")pass standalone.
--enable-deepseek-v4-fp4-indexer --method mtp):RULER niah @ 16k (
niah_single_1/2,niah_multikey_1, limit 50) = 1.0.🤖 Generated with Claude Code