feat(deepseek_v4): Add fp4 indexer fo dsv4 - #1709
Open
junhaha666 wants to merge 17 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>
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>
`build_for_cudagraph_capture` captures EVERY decode graph through the ragged indexer branch whenever `dspark.ragged` is on, but `_build_v4_indexer_meta` only refreshed the ragged varlen windows on ragged steps. Since `graph_key` is just `(graph_bs, max_q_len)`, a rectangular step (DP-sync dummy, boundary / reorder, no-shrink) replays a ragged-captured graph whose baked kernels then read the previous step's stale windows. The paged-KV `buffer_load` in `pa_mqa_logits_fp4_prefill_kernel_0` runs with bounds-check off (num_records=0xffffffff), so the bad address faults instead of clamping — MEMORY_VIOLATION mid-run, confirmed via rocm-debug-agent. Refresh the windows on every decode fwd instead: a rectangle is just the uniform case, so rectangular steps feed a `next_n` fill (the same shape capture synthesizes). Extracted into `_refresh_fp4_ragged_windows` so both branches share one implementation. With this, FULL + ragged no longer faults, so the fail-fast guard in `_score_topk_decode_ragged_fp4` is removed. Validated on DeepSeek-V4-Pro-DSpark (fp4 indexer, spec=7, DP-attention): - FULL + ragged: gsm8k 95.53%, acceptance ~38% (crashed before the fix) - 8k/1k x 1280 req: 1280/1280 successful, zero GPU faults - PIECEWISE unaffected (rect steps just refresh one extra buffer) Also in this change: - Deduplicate the FP4 mqa schedule constants, which were defined twice with "MUST match" comments. They now live in `v4_kernels` as `FP4_MQA_PARALLEL_UNIT_NUM` / `FP4_MQA_BLOCK_K`; the old names said DECODE even though both the decode and prefill paths use them. - Drop the `parallel_unit_num=` kwarg from both mqa call sites. Passing `cta_info`/`n_ctas` makes the kernel skip its internal schedule build, so the kwarg was dead — and its hardcoded 512 contradicted the builder's real `max(512, T_dec)`. - Compute the prefill chunk's CTA count only in the multi-chunk branch that uses it, instead of on every call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…nce #1697 f3b1fb4 fixed a MEMORY_VIOLATION under `--cudagraph-mode FULL` + DSpark ragged by also refreshing the ragged varlen windows on rectangular decode steps: at that point `prepare_decode` set `dspark_ragged_lens_gpu` only `if ragged_lens is not None`, so a rectangular step left the windows holding the previous step's contents while the captured graph replayed the ragged kernels over them. Merging origin/main brought in #1697 (2c682dc), which widened that condition to if ragged_lens is not None or _dspark_ragged_graph: with `_dspark_ragged_graph = dspark.ragged and confidence_schedule`. Whenever those options are on the field is now set on every decode — uniform lengths when nothing shrank — so `_build_v4_indexer_meta` always takes the ragged branch and the rectangular one only runs with the options off. The refresh added to that branch has therefore been dead code since the merge; remove it and describe what actually happens now. The remaining per-forward rebuild is not redundant: `compute_varqlen_windows` takes `n_committed_csa_per_seq` (= context_len // 4), which advances on every decode step, so each row's window end moves even when no per-seq length shrank. Verified post-merge on DeepSeek-V4-Pro-DSpark (fp4 indexer, spec=7, DP-attention) with `--cudagraph-mode FULL` + ragged, the combination f3b1fb4 was written for: CUDAGraph capture succeeds, 16 concurrent requests pass (this load faulted before that fix), gsm8k 95.38% flexible / 95.22% strict, acceptance 37.6%, zero MEMORY_VIOLATION. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
CI (black 26.5.1 + ruff via reviewdog) failed on this branch. Fixes exactly what it reported, nothing else: - black: 7 files reformatted. - F821 in deepseek_v4_attn.py: `_refresh_fp4_ragged_windows` annotated its `meta` argument as `typing.Dict`, which resolves locally but not on CI, where the run sits on a base that no longer imports it. Use the builtin `dict[str, Any]`, which is valid either way and matches the PEP 585/604 style the other fixes move toward. - UP045 (65 sites): `Optional[X]` -> `X | None` in module_dispatch_ops.py, v4_kernels/fused_compress.py, models/deepseek_v4.py and plugin/vllm/models/deepseek_v4.py. - I001: import blocks in v4_kernels/__init__.py, models/deepseek_v4.py and tests/test_decode_indices_paged.py. - RUF022: `__all__` in v4_kernels/__init__.py, which the FP4_MQA_* constants had been appended to. - RUF100: two `# noqa: E402` in tests/test_decode_indices_paged.py. Deliberately left alone: the UP045 and I001 hits ruff reports in deepseek_v4_attn.py, since CI flagged only F821 there and the rest sit outside the diff; and the `# noqa: F401` on the `module_dispatch_ops` side-effect import in models/deepseek_v4.py, which a local run without the repo's rule selection calls unused but which CI does not flag and the op registration needs. `pytest tests/` gives an identical 53 failed / 746 passed / 32 skipped / 19 errors before and after — all pre-existing, from missing optional deps. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist