Fix dual-chunk sparse fallback index overflow (#27361)
This commit is contained in:
Han Yu
@@ -6,7 +6,7 @@ unit-test suite, **organized by the action needed to address it**.
|
||||
|
||||
Anything failing that is not listed here should be treated as a regression.
|
||||
|
||||
Last updated: 2026-05-29
|
||||
Last updated: 2026-06-05
|
||||
|
||||
## Reference runs
|
||||
|
||||
@@ -182,8 +182,6 @@ moment production is fixed; no test method invokes them today.
|
||||
|
||||
| Citation | Symptom | Trigger | Status |
|
||||
|---|---|---|---|
|
||||
| `dual_chunk_flashattention_backend.py:1110-1132` | `RuntimeError: The size of tensor a (4) must match the size of tensor b (5)` at `vertical_buffer.copy_()` | `vertical_size ≤ 5`: fallback `torch.arange(0, intra_K_size, max(1, intra_K_size/5))` returns up to 5 elements into `vertical_size=4` buffer when `intra_vertical_indices.nelement() == 0` | `[no test]` (`dual_chunk/README.md`); smoke helper `run_dual_chunk_sparse_sub_window_case` wired but not invoked |
|
||||
| `_vertical_slash_sparse_attention` (`convert_vertical_slash_indexes` block math) | `cudaErrorIllegalAddress` deep inside the kernel | `vertical_size=8` with `seq_len ≥ 128`: unstated invariant that `vertical_size + slash_size >= chunk_len_blocks` | `[no test]` (same smoke helper) |
|
||||
| Triton dense `DRAFT_EXTEND` (non-V2) | Eager fixture/reference mismatch on narrow accepted-token layouts | Test omitted | `[no test]` (`dense/README.md`) |
|
||||
|
||||
## C.6. DSA-specific structural gaps
|
||||
|
||||
@@ -19,6 +19,7 @@ Columns are runner modes; rows are kernel-path modes of the single
|
||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
||||
| Non-sparse | ✓ first-window, successor-chunk, inter-chunk extend/decode layouts + GQA decode | deferred: graph metadata for dual-chunk not scoped | deferred | deferred | blocked: `init_forward_metadata` asserts `is_prefill() or is_decode()` (`dual_chunk_flashattention_backend.py:179`); `TARGET_VERIFY` falls under `is_prefill()` but the wrapper hasn't been wired through | deferred | deferred | deferred | blocked: `DRAFT_EXTEND_V2` excluded from `is_prefill()` alias (see Production-Unsupported below) | deferred | deferred | — |
|
||||
| Sparse all-column (`vertical_size`/`slash_size` chosen so every key in the first chunk is selected) | ✓ single-request first-chunk, multi-request first-chunk, page-boundary first-chunk | — | — | — | blocked: same `is_prefill` assertion | — | — | — | blocked: same | — | — | — |
|
||||
| Sparse sub-window (`vertical_size=4`, `slash_size=4`, `seq_len=128`) | ✓ independent DCA top-k/split/fallback reference + torch sparse-kernel reference | — | — | — | — | — | — | — | — | — | — | — |
|
||||
| Threshold-gated sparse (`sparse_attention_threshold=100`, seq_len=16 → gate disables sparse, falls back to dense) | ✓ verifies `current_orig_seq_len > threshold` gate semantics | — | — | — | — | — | — | — | — | — | — | — |
|
||||
|
||||
## Input And Config Coverage
|
||||
@@ -33,6 +34,11 @@ Columns are runner modes; rows are kernel-path modes of the single
|
||||
(≤16 tokens) so the dense reference remains valid.
|
||||
- Multi-request sparse and page-boundary sparse variants exercise per-request
|
||||
`cu_seqlens_*` slicing inside `_dual_chunk_flash_attn_prefill_func`.
|
||||
- Sub-window sparse prefill uses `vertical_size=4`, `slash_size=4`, and
|
||||
`seq_len=128` to verify the DCA-specific content-aware top-k split and
|
||||
empty-stage fallback against an independent reference, then verifies the
|
||||
sparse output against a torch sparse-kernel reference that consumes the
|
||||
production block/column metadata.
|
||||
- Threshold-gated sparse uses `sparse_attention_threshold=100` so a 16-token
|
||||
prompt bypasses the sparse kernel and falls through to the dense chunk
|
||||
flash path, exercising the gate semantics in the wrapper.
|
||||
@@ -77,65 +83,8 @@ See `KNOWN_FAILURES.md` §1 for the full root cause + fix.
|
||||
|
||||
- Populate CUDA graph and PCG/BCG runner metadata after eager non-sparse
|
||||
coverage is stable across more chunk layouts.
|
||||
- **Sub-context-window sparse pruning reference (genuine follow-up)** —
|
||||
The current "all-column" sparse cases match the dense reference exactly
|
||||
because the chosen `vertical_size=16` + `slash_size=16` + `last_q=16`
|
||||
configuration covers every column in the first chunk for `seq_len <= 16`.
|
||||
A truly pruning case needs `seq_len >> vertical_size + slash_size` and a
|
||||
reference that applies the same mask the kernel applies.
|
||||
|
||||
The blocker is that the production sparse-attention config
|
||||
`("vertical_and_slash", v_size, s_size, threshold)` is **content-aware**:
|
||||
per-head `v_idx` and `s_idx` are picked by top-k attention scores over
|
||||
the last `last_q` queries, not from a fixed schedule
|
||||
(`dual_chunk_flashattention_backend.py:_dual_chunk_flash_attn_prefill`).
|
||||
An independent reference therefore has three paths:
|
||||
|
||||
1. **Mock the sparse-config lookup** — patch
|
||||
`get_sparse_attention_config` or the per-layer top-k selection so the
|
||||
fixture supplies known `v_idx` / `s_idx` tensors. Then write a
|
||||
token-level reference that masks `attn_scores[q, k] = -inf` unless
|
||||
`k in v_idx` or `(q - k) in s_idx` (with causal `k <= q`). This is the
|
||||
cleanest path but needs a hook in `_dual_chunk_flash_attn_prefill_func`
|
||||
that doesn't exist today.
|
||||
2. **Replicate `convert_vertical_slash_indexes`** at block granularity in
|
||||
pure-PyTorch, then iterate `(block_count, block_offset, column_count,
|
||||
column_index)` to build a per-(query_block, key_block) mask matching
|
||||
the kernel's selection. Faithful but tedious — the block math (M=64,
|
||||
N=64) needs to be mirrored exactly.
|
||||
3. **Statistical recovery check** — compute dense attention scores
|
||||
`softmax(Q @ K^T)` per head, identify the top-k columns by score, and
|
||||
verify the sparse kernel output approximates the dense output modulo
|
||||
the dropped probability mass. Not strict `assert_close`; rejects only
|
||||
gross divergences.
|
||||
|
||||
Option 1 is recommended. It requires either: (a) a new
|
||||
`sparse_attention_config_override` kwarg threaded through
|
||||
`DualChunkFlashAttentionBackend.__init__` that bypasses the content-aware
|
||||
selection, or (b) monkeypatching `get_sparse_attention_config` on the
|
||||
fixture's backend instance. Until that lands, the all-column sparse +
|
||||
threshold-gated cases keep the kernel/wrapper integration covered but
|
||||
the per-column sparse math is unverified.
|
||||
|
||||
**Production-side bugs surfaced while attempting Option 3
|
||||
(smoke-test "sparse output != dense output"):** two issues block even a
|
||||
smoke-only sub-window test today.
|
||||
|
||||
- `dual_chunk_flashattention_backend.py:1110-1122`: when a chunk's
|
||||
`intra_vertical_indices.nelement() == 0`, the fallback appends
|
||||
`torch.arange(0, intra_K_size, max(1, intra_K_size/5))`. With
|
||||
`intra_K_size=48` this is `arange(0, 48, 9.6)` → 5 elements, but the
|
||||
`vertical_buffer` is sized to `vertical_size` (=4 in a sub-window
|
||||
config). The copy at line 1132 then raises
|
||||
`RuntimeError: The size of tensor a (4) must match the size of
|
||||
tensor b (5)`. The fallback should clip to `vertical_size` slots.
|
||||
- With `vertical_size=8` to clear the overflow, the sparse kernel
|
||||
crashes with `cudaErrorIllegalAddress` deep inside
|
||||
`_vertical_slash_sparse_attention`, suggesting the
|
||||
`convert_vertical_slash_indexes` block math has an unstated
|
||||
invariant that `vertical_size + slash_size >= chunk_len_blocks` or
|
||||
similar. Needs a kernel-side audit.
|
||||
|
||||
The smoke-test helper `run_dual_chunk_sparse_sub_window_case` is wired
|
||||
through `common/attention_methods/dual_chunk_attention.py` for when
|
||||
those production bugs are fixed; no test method invokes it today.
|
||||
- **Broaden sub-window sparse coverage** — the current regression case covers
|
||||
`prefix_lens=(0,)`, `extend_lens=(128,)`, and no GQA. Add
|
||||
multi-request batches, nonzero prefixes, GQA, and more sparse config variants
|
||||
once those paths need explicit sparse pruning coverage. The 64x64
|
||||
vertical/slash converter remains covered at the sgl-kernel layer.
|
||||
|
||||
@@ -13,9 +13,11 @@ from sglang.test.kits.attention_unittest.attention_methods.dual_chunk_attention
|
||||
DualChunkAttentionCase,
|
||||
make_dual_chunk_cases,
|
||||
make_dual_chunk_sparse_cases,
|
||||
make_dual_chunk_sparse_sub_window_cases,
|
||||
make_dual_chunk_sparse_threshold_gated_cases,
|
||||
run_dual_chunk_attention_case,
|
||||
run_dual_chunk_sparse_attention_case,
|
||||
run_dual_chunk_sparse_sub_window_case,
|
||||
run_dual_chunk_sparse_threshold_gated_case,
|
||||
)
|
||||
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
|
||||
@@ -75,6 +77,9 @@ class TestDualChunkFlashAttentionBackendCorrectness(CustomTestCase):
|
||||
SPARSE_THRESHOLD_GATED_CASES = make_dual_chunk_sparse_threshold_gated_cases(
|
||||
"dual_chunk_flash_attn"
|
||||
)
|
||||
SPARSE_SUB_WINDOW_CASES = make_dual_chunk_sparse_sub_window_cases(
|
||||
"dual_chunk_flash_attn"
|
||||
)
|
||||
# Replay prefix_lens must each be >= capture_prefix_len (= fill-value - 1).
|
||||
# Dual-chunk's `get_cuda_graph_seq_len_fill_value()` returns 1, so capture
|
||||
# uses prefix=0. We pick a 3-request batch with varied lengths to exercise
|
||||
@@ -106,34 +111,10 @@ class TestDualChunkFlashAttentionBackendCorrectness(CustomTestCase):
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_dual_chunk_sparse_threshold_gated_case(self, case)
|
||||
|
||||
# Sub-context-window sparse pruning: BLOCKED on production-side
|
||||
# edge cases.
|
||||
#
|
||||
# The `run_dual_chunk_sparse_sub_window_case` helper in
|
||||
# `common/attention_methods/dual_chunk_attention.py` is left in
|
||||
# place for when those production gaps are fixed, but no test
|
||||
# method invokes it today. See `dual_chunk/README.md` →
|
||||
# "Sub-context-window sparse pruning" for the engineering paths
|
||||
# and the two production bugs surfaced while attempting to land
|
||||
# this coverage:
|
||||
#
|
||||
# - `dual_chunk_flashattention_backend.py:1110-1122`: when a chunk's
|
||||
# `intra_vertical_indices.nelement() == 0`, the fallback appends
|
||||
# `torch.arange(0, intra_K_size, max(1, intra_K_size/5))` which
|
||||
# can produce more elements than the `vertical_size`-slot buffer
|
||||
# allows, raising `RuntimeError: The size of tensor a (4) must
|
||||
# match the size of tensor b (5)`. Triggered by
|
||||
# `vertical_size in [4, 5]` with `seq_len=128`.
|
||||
# - With `vertical_size=8` to avoid the overflow above, the sparse
|
||||
# kernel raises a `cudaErrorIllegalAddress` deep inside
|
||||
# `_vertical_slash_sparse_attention`, suggesting the
|
||||
# `convert_vertical_slash_indexes` block math expects different
|
||||
# invariants than what a `vertical_size + slash_size < chunk_len`
|
||||
# config supplies.
|
||||
#
|
||||
# The all-column + threshold-gated cases above keep the integration
|
||||
# path covered; sub-window correctness needs production hardening
|
||||
# before unit-test coverage is safe.
|
||||
def test_sparse_dual_chunk_sub_window_cases(self):
|
||||
for case in self.SPARSE_SUB_WINDOW_CASES:
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_dual_chunk_sparse_sub_window_case(self, case)
|
||||
|
||||
def test_runner_mode_cuda_graph_decode_cases(self):
|
||||
for case in self.CUDA_GRAPH_DECODE_CASES:
|
||||
|
||||
Reference in New Issue
Block a user