Skip to content

[Feature][P2] Enable elastic 1M-context decode attention with block size 32 #962

Description

@sjduan

Parent tracking issue

Summary

Enable elastic 1M-context decode for models/deepseek_v4_flash_dspark/ while preserving the existing DeepSeek-V4-Flash attention math and the deployment point defined by #905 (B=64, S=8, block size 32).

Here, 1M is a capacity ceiling, not a fixed execution shape. Runtime work must follow the actual visible context of each request. Short requests must not allocate or submit 1M-sized attention/indexer work, and 128/16K/1M must not become separate kernel profiles.

Goals

  1. Make block size 32 the canonical decode cache ABI.
  2. Support per-request logical positions and cache capacity up to 1,048,576 tokens.
  3. Preserve heterogeneous and inactive requests through ragged runtime metadata.
  4. Enable all three decode attention modes:
    • SWA / ratio 0
    • HCA / ratio 128
    • CSA / ratio 4 with exact Top-512
  5. Keep production attention structure as close to the existing baseline as possible; long-context support should primarily change metadata, persistent-cache addressing, and the amount of submitted work.

Required design

Common geometry and metadata

  • Canonical capacity: MAX_CONTEXT_TOKENS = 1_048_576.
  • Cache storage block size: 32 rows.
  • Use allocator-owned global physical pools; do not pre-partition cache capacity by a fixed batch.
  • Device entries consume ragged descriptors containing physical page IDs, valid ranges, epochs, write slots, and exact packed work.
  • Page IDs and write slots are local to one layer's pool view and must fail closed when invalid, stale, inactive, or out of range.
  • RoPE inputs are token/event-local. Do not pass a [1M, ...] RoPE table.
  • Runtime length changes task counts and tail validity only; it must not change kernel semantics or select a length-specific profile.

SWA (ratio 0)

  • Semantic window remains 128 rows, stored as four 32-row pages per admitted request.
  • swa_sources[Q, 128] represents persistent rows, invalid rows, and current-step overlay rows.
  • Current-step KV participates causally through the overlay and is committed to the persistent ring only after attention has finished reading the old rows.
  • Contexts from 128 through 1M use the same 128-row attention tile and persistent capacity.

HCA (ratio 128)

  • A 1M request produces at most 8,192 compressed rows.
  • Attention uses fixed 128-row shards; the runtime submits 0–64 shards per query from exact packed work.
  • The ratio-128 state ring retains 128 semantic rows using sixteen 8-row state pages.
  • Compressed KV uses ragged 32-row pages and explicit page epochs.
  • Boundary compression events use event-local RoPE and explicit main/state write dependencies.
  • Raw SWA history plus compressed HCA history are merged with online softmax; no dense history index or topk_all tensor is allowed.

CSA (ratio 4)

  • A 1M request produces at most 262,144 candidates.
  • Fixed selector leaf: 2,048 candidates -> sorted Top-512.
  • Runtime submits only ceil(valid_candidates / 2048) active leaves (0–128 per query), followed by a bounded exact merge forest.
  • Main KV, index cache, main state, and inner state use independent allocator-owned pools.
  • Index cache remains a flat-row ABI; candidate/page work is ragged and exact.
  • Preserve the baseline HC/QKV/compressor/sparse-value structure. If micro-chunking is required for the selector forest, confine it to the indexer work instead of repeating the full attention pipeline.

Production code scope

Primary files:

  • models/deepseek_v4_flash_dspark/config.py
  • models/deepseek_v4_flash_dspark/context_geometry.py
  • models/deepseek_v4_flash_dspark/decode_metadata.py
  • models/deepseek_v4_flash_dspark/decode_compressor_ratio128.py
  • models/deepseek_v4_flash_dspark/decode_compressor_ratio4.py
  • models/deepseek_v4_flash_dspark/decode_indexer_compressor.py
  • models/deepseek_v4_flash_dspark/decode_indexer.py
  • models/deepseek_v4_flash_dspark/decode_indexer_topk.py
  • models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py
  • models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py
  • models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py
  • models/deepseek_v4_flash_dspark/decode_swa.py
  • models/deepseek_v4_flash_dspark/decode_hca.py
  • models/deepseek_v4_flash_dspark/decode_csa.py

Standalone fixtures, golden functions, trace-only CLI options, and historical compatibility aliases are validation scaffolding. They should not remain in the final production attention modules after correctness and performance validation has been transferred to tests.

Acceptance criteria

Static and host checks

  • All public tensor ABIs have at most five dimensions.
  • No production ABI contains a full 1M RoPE table, dense 1M block table, or fixed max-work tensor per request.
  • Page/range/epoch/write-slot validation rejects stale, missing, aliased, inactive, and out-of-range metadata.
  • Geometry covers 1, 127, 128, 129, 12K, 16K, 32K, 512K, 1M-1, and 1M.
  • Heterogeneous request lengths produce independent SWA windows, HCA shard counts, and CSA leaf counts.
  • Block-size-32 page boundaries and ring rollover are explicitly covered.

Device correctness

  • SWA standalone passes short history, ring wrap, heterogeneous lengths, and 1M tail.
  • Ratio-128 compressor passes no-event, boundary event, state rollover, page permutation, and 1M tail.
  • HCA sparse attention passes 0/1/2/64-shard cases and mixed lengths.
  • CSA selector covers candidate boundaries around 0/1/511/512/513/2048 and the 128-leaf 1M case.
  • CSA sparse-value attention passes 16K and 1M-tail cases.
  • Integrated SWA, HCA, and CSA outputs match their numerical goldens without relaxing existing thresholds.
  • Inactive requests submit no persistent write and preserve the output sentinel.

Work and trace evidence

  • SWA task/workspace shape is constant once the 128-row window is full.
  • HCA work scales with the actual number of 128-row compressed shards.
  • CSA work scales with the actual number of 2,048-candidate leaves.
  • Representative traces cover 128, 16K, and 1M logical contexts.
  • No heap/task/dependency/TensorMap overflow or scheduler deadlock occurs in the supported attention cases.

Non-goals

Completion definition

This issue is complete when all three standalone attention paths use block-32 ragged cache metadata, execute the exact work implied by heterogeneous runtime lengths up to 1M, pass their existing numerical thresholds, and have no remaining production dependency on the old dense 16K ABI.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions