Delete cutlass_mla, non-Marlin GPTQ, AWQ AOT kernel, and Dual Chunk Flash Attention (#32114)
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
co-authored by
Brayden Zhong
Brayden Zhong
parent
7152c14384
commit
c0b790cf7f
@@ -34,29 +34,6 @@ hardware/version X).
|
||||
|
||||
# A. Container re-image required
|
||||
|
||||
## A.1. `flash_attn` SM10.x wheel missing
|
||||
|
||||
**Affected**: `dual_chunk/test_dual_chunk_flash_attn.py` (entire class — 5
|
||||
test methods, ~18 subtests)
|
||||
|
||||
**Symptom on GB300**:
|
||||
```
|
||||
ImportError: cannot import name 'flash_attn_varlen_func' from 'flash_attn'
|
||||
```
|
||||
|
||||
**Root cause**: `DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func`
|
||||
via `sglang.kernels.ops.attention.flash_attention`. On SM 8.x / 9.x that resolves to
|
||||
sgl-kernel's FA3 build (works on H200). On other SMs, the JIT kernel falls
|
||||
back to the upstream `flash_attn` (FA2) wheel — but the
|
||||
`lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package on
|
||||
SM10.x is missing `flash_attn_varlen_func`.
|
||||
|
||||
**Gate**: `_dual_chunk_fa_supported()` in
|
||||
`dual_chunk/test_dual_chunk_flash_attn.py` skips the whole class on the
|
||||
fallback-broken path. Hopper passes through unchanged.
|
||||
|
||||
**Fix**: Re-image with an SM10.x-compiled `flash_attn` wheel.
|
||||
|
||||
## A.2. tilelang `wait_wgmma` template missing on SM10.x
|
||||
|
||||
**Affected**:
|
||||
@@ -89,7 +66,6 @@ so that "skipped: ..." results have a quick lookup.
|
||||
|
||||
| Backend | Required SM | Gate location | Error if unguarded |
|
||||
|---|---|---|---|
|
||||
| `cutlass_mla` | exactly SM 10.0 (B200) | `mla/test_cutlass_mla.py::_supported` | `cutlass_mla_decode is only supported on compute capability 10.0, but found sm version 103` |
|
||||
| `flashmla` decode/verify | SM 9.0 (Hopper) only | `mla/test_flashmla.py:_DECODE_REQUIRES_SM90A` | `Dense decode MLA is only supported on SM90a architecture` |
|
||||
| `trtllm_mla` | SM 12.0a / 12.1a | `mla/test_trtllm_mla.py::_supported` | FlashInfer XQA MLA dispatch reject |
|
||||
| `tokenspeed_mla` | SM ≥ 10.0 + FP8 KV + pkg | `mla/test_tokenspeed_mla.py::_supported` | `tokenspeed_mla` import or kernel dispatch |
|
||||
@@ -99,7 +75,7 @@ so that "skipped: ..." results have a quick lookup.
|
||||
| `fa3` (non-MLA) | SM 80 or SM 90 | `_is_fa3_supported` in `flash_attention_v3.py` | `attention_registry.py:177-180` reject |
|
||||
|
||||
**SM10.3 vs SM10.0**: GB300 is SM10.3. Gates that require exactly SM10.0
|
||||
(cutlass_mla, dsa trtllm) intentionally skip on GB300 because the kernel
|
||||
(dsa trtllm) intentionally skip on GB300 because the kernel
|
||||
binaries in the container aren't compiled for sm_103. Flip the gates to
|
||||
`major == 10` (drop the `minor == 0`) once GB300-compiled binaries land.
|
||||
|
||||
@@ -133,13 +109,7 @@ methods that record each backend's failure mode inline as
|
||||
| `mla/test_flashmla.py::test_layout_robustness_cases` (extend) | `non_monotonic_extend` | FlashMLA extend raises CUDA illegal memory access. |
|
||||
| `mla/test_flashmla.py::test_layout_robustness_cases` (decode) | `interleaved_pages` | FlashMLA decode raises `shape '[-1, 64, 1, 32]' is invalid for input of size N`. |
|
||||
|
||||
### Dual-chunk
|
||||
|
||||
| Test | Layout | Root cause |
|
||||
|---|---|---|
|
||||
| `dual_chunk/test_dual_chunk_flash_attn.py::test_layout_robustness_cases` (extend) | `non_monotonic_extend` | `_dual_chunk_flash_attn_prefill_func` uses `cu_seqlens_*` indexing into contiguous K slots (`dual_chunk_flashattention_backend.py:834+`); scattered extend-token slots break that contiguity. |
|
||||
|
||||
**Total**: 9 layout-handling production bugs documented.
|
||||
**Total**: 8 layout-handling production bugs documented.
|
||||
|
||||
## C.2. Speculative-mode rejects
|
||||
|
||||
@@ -205,7 +175,6 @@ fail at backend init.
|
||||
| Backend | Required page size(s) | Citation |
|
||||
|---|---|---|
|
||||
| FlashMLA | `64` only | `server_args.py:2767-2770` |
|
||||
| Cutlass MLA | `128` only | `server_args.py:2776-2779`, `cutlass_mla_backend.py:31` |
|
||||
| TRT-LLM MLA | `{32, 64}` | `server_args.py:2790-2794` |
|
||||
| Tokenspeed MLA | `{32, 64}` | `server_args.py:2809-2813`, `tokenspeed_mla_backend.py:111-113` |
|
||||
| TRT-LLM MHA | `{16, 32, 64}` | `server_args.py:2849-2853` |
|
||||
@@ -241,12 +210,9 @@ fail at backend init.
|
||||
|
||||
| Test file | Failure type | Section |
|
||||
|---|---|---|
|
||||
| `dual_chunk/test_dual_chunk_flash_attn.py` | Container: `flash_attn` SM10.x wheel | §A.1 |
|
||||
| `dual_chunk/test_dual_chunk_flash_attn.py::test_layout_robustness_cases` (non_monotonic_extend) | Layout-handling bug | §C.1 |
|
||||
| `dsa/test_dsa.py::test_sparse_tilelang_*` | Container: tilelang `wait_wgmma` | §A.2 |
|
||||
| `dsa/test_dsa.py::test_sparse_*_impl_variants` (tilelang row) | Container: tilelang `wait_wgmma` | §A.2 |
|
||||
| `dsa/test_dsa.py::test_sparse_*_impl_variants` (fa3 / trtllm rows) | Hardware gate | §B |
|
||||
| `mla/test_cutlass_mla.py` (all) | Hardware gate (SM 10.0 exactly) | §B |
|
||||
| `mla/test_flashmla.py` (DECODE/verify subtests) | Hardware gate (SM 9.0 Hopper) | §B |
|
||||
| `mla/test_flashinfer.py::test_runner_mode_eagle_draft_cuda_graph_runner_cases` | Backend bug gated on SM≥10 | §C.3 |
|
||||
| `mla/test_flashinfer.py::test_layout_robustness_cases` | Layout-handling bug | §C.1 |
|
||||
|
||||
@@ -1,90 +0,0 @@
|
||||
# Dual-Chunk Attention Capability Matrix
|
||||
|
||||
This folder covers dual-chunk attention tests. `dual_chunk_flash_attn` is not
|
||||
a dense backend swap: it expects a packed five-way query projection (`query`,
|
||||
`succ`, `inter`, and critical variants), so the dense Q/K/V harness is
|
||||
structurally wrong for this method. The single attention backend here is
|
||||
`dual_chunk_flash_attn`; the rows below distinguish kernel-path modes.
|
||||
|
||||
## Coverage Matrix
|
||||
|
||||
Columns are runner modes; rows are kernel-path modes of the single
|
||||
`dual_chunk_flash_attn` backend. Cells use:
|
||||
- **✓ \<variants\>** — exercised, with the config variants listed in the cell
|
||||
- **—** — not applicable / not exercised
|
||||
- **blocked: \<reason\>** — production-unsupported, not a follow-up
|
||||
- **deferred: \<reason\>** — could land later, currently disabled
|
||||
|
||||
| Kernel path | Eager Phase 2 | CG decode | PCG extend | BCG extend | Verify eager | Verify CG | DE eager | DE CG | DE-V2 CG | EAGLE-draft runner | EAGLE-DE runner | FKVMTP runner |
|
||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
||||
| 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
|
||||
|
||||
- Page size 1 extend, exact-page extend, page-boundary crossing extend, and
|
||||
ragged extend batches.
|
||||
- Decode page-boundary coverage and GQA decode coverage.
|
||||
- Successor-chunk and inter-chunk extend/decode layouts where `query_succ`
|
||||
and `query_inter` are active and use independent projection weights.
|
||||
- Sparse all-column prefill uses `head_dim=128` to match the local sparse
|
||||
FlashAttention build and selects every column in the first chunk
|
||||
(≤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.
|
||||
|
||||
## Container Gate (SM10.x)
|
||||
|
||||
`DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func` via
|
||||
`sglang.kernels.ops.attention.flash_attention`. On SM8.x / SM9.x that resolves to sgl-kernel's
|
||||
FA3 build; on SM != {8, 9} (notably SM10.x / GB300) the JIT kernel falls back
|
||||
to the upstream `flash_attn` (FA2) wheel, which the
|
||||
`lmsysorg/sglang:nightly-dev-cu13` container ships without an SM10.x-compiled
|
||||
`flash_attn_varlen_func`. `test_dual_chunk_flash_attn.py` probes both paths at
|
||||
module import: if FA3 is supported (`major in {8, 9}`) it runs unconditionally;
|
||||
otherwise it tries `from flash_attn import flash_attn_varlen_func` and skips the
|
||||
whole class with the documented reason if the symbol is missing. Re-image with
|
||||
an SM10.x-compiled flash_attn wheel to clear; no test-code change needed.
|
||||
|
||||
See `KNOWN_FAILURES.md` §1 for the full root cause + fix.
|
||||
|
||||
## Production-Unsupported
|
||||
|
||||
- **Non-prefill / non-decode forward modes** —
|
||||
`dual_chunk_flashattention_backend.py:179` asserts
|
||||
`forward_mode.is_prefill() or forward_mode.is_decode()`. `is_prefill()`
|
||||
aliases to `is_extend()` (`forward_batch_info.py:103-104`) and covers
|
||||
`EXTEND` / `MIXED` / `DRAFT_EXTEND` / `TARGET_VERIFY` / `SPLIT_PREFILL` /
|
||||
`DLLM_EXTEND`, but `DRAFT_EXTEND_V2` is excluded by default. So
|
||||
`DRAFT_EXTEND_V2` is structurally unreachable for `dual_chunk_flash_attn`.
|
||||
- **Non-causal / windowed-attention requests** — `forward_extend` raises
|
||||
`ValueError("Dual Chunk Attention does not support causal=False")`
|
||||
(`dual_chunk_flashattention_backend.py:698`) and
|
||||
`ValueError("Dual Chunk Attention does not support window_size")`
|
||||
(`dual_chunk_flashattention_backend.py:700`).
|
||||
- **Sparse mode `chunk_len % block_size != 0`** — raises
|
||||
`ValueError("chunk_len must be divisible by block_size.")`
|
||||
(`dual_chunk_flashattention_backend.py:860, 1491`). The current fixture
|
||||
picks divisible values.
|
||||
- **Unsupported `head_dim`** — only `head_dim in {16, 32, 64, 128, 256, 512}`
|
||||
is accepted (`dual_chunk_flashattention_backend.py:1611`).
|
||||
|
||||
## Next Work
|
||||
|
||||
- Populate CUDA graph and PCG/BCG runner metadata after eager non-sparse
|
||||
coverage is stable across more chunk layouts.
|
||||
- **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.
|
||||
@@ -1 +0,0 @@
|
||||
"""Dual-chunk attention unit-test package."""
|
||||
@@ -1,177 +0,0 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
from sglang.test.kits.attention_unittest.attention_methods.dual_chunk_attention import (
|
||||
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 (
|
||||
run_dual_chunk_cuda_graph_decode_case,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
# Container gate (KNOWN_FAILURES.md §1): `DualChunkFlashAttentionBackend` calls
|
||||
# `flash_attn_varlen_func` on every forward via
|
||||
# `sglang.kernels.ops.attention.flash_attention`. On SM8x/SM9x, that resolves to sgl-kernel's
|
||||
# FA3 build (which works). On SM != {8, 9} (notably SM10.3 / GB300), the JIT
|
||||
# kernel falls back to the upstream `flash_attn` (FA2) wheel — but the
|
||||
# `lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package ships
|
||||
# without `flash_attn_varlen_func` on SM10.x, so every dual-chunk forward
|
||||
# fails at import time inside the fallback. Skip the whole suite only when
|
||||
# that fallback path is actually broken (not on Hopper, where we never enter it).
|
||||
# Re-image the container with an SM10.3-compiled flash_attn wheel to clear.
|
||||
def _dual_chunk_fa_supported() -> tuple[bool, str]:
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA is required"
|
||||
major, _minor = torch.cuda.get_device_capability()
|
||||
# FA3 path is taken when sm major is 8 or 9 (see
|
||||
# `sglang.kernels.ops.attention.flash_attention_v3._is_fa3_supported`). On that path
|
||||
# the upstream `flash_attn` fallback is never invoked.
|
||||
if major in (8, 9):
|
||||
return True, ""
|
||||
# Otherwise (sm 7.x or sm >= 10.x) the JIT kernel falls back to upstream
|
||||
# `flash_attn.flash_attn_varlen_func`. Probe it; if missing, skip.
|
||||
try:
|
||||
from flash_attn import ( # noqa: F401
|
||||
flash_attn_varlen_func as _flash_attn_varlen_func,
|
||||
)
|
||||
|
||||
return True, ""
|
||||
except ImportError as exc:
|
||||
return False, (
|
||||
f"flash_attn_varlen_func is not available in upstream `flash_attn` "
|
||||
f"(SM{major}.x JIT-kernel fallback): {exc}. "
|
||||
f"Re-image the container with an SM{major}.x-compiled flash_attn wheel."
|
||||
)
|
||||
|
||||
|
||||
_DUAL_CHUNK_FLASH_ATTN_AVAILABLE, _DUAL_CHUNK_SKIP_REASON = _dual_chunk_fa_supported()
|
||||
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b", runner_config="4-gpu-b200")
|
||||
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
|
||||
@unittest.skipIf(not _DUAL_CHUNK_FLASH_ATTN_AVAILABLE, _DUAL_CHUNK_SKIP_REASON)
|
||||
class TestDualChunkFlashAttentionBackendCorrectness(CustomTestCase):
|
||||
CASES = make_dual_chunk_cases("dual_chunk_flash_attn")
|
||||
SPARSE_CASES = make_dual_chunk_sparse_cases("dual_chunk_flash_attn")
|
||||
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
|
||||
# both the page-boundary and within-page slots.
|
||||
CUDA_GRAPH_DECODE_CASES = (
|
||||
DualChunkAttentionCase(
|
||||
name="runner_cuda_graph_dual_chunk_decode_page_boundary",
|
||||
backend="dual_chunk_flash_attn",
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_heads=4,
|
||||
num_kv_heads=4,
|
||||
page_size=16,
|
||||
prefix_lens=(14, 15, 16),
|
||||
),
|
||||
)
|
||||
|
||||
def test_projected_dual_chunk_attention_cases(self):
|
||||
for case in self.CASES:
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_dual_chunk_attention_case(self, case)
|
||||
|
||||
def test_sparse_dual_chunk_attention_cases(self):
|
||||
for case in self.SPARSE_CASES:
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_dual_chunk_sparse_attention_case(self, case)
|
||||
|
||||
def test_sparse_dual_chunk_threshold_gated_cases(self):
|
||||
for case in self.SPARSE_THRESHOLD_GATED_CASES:
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_dual_chunk_sparse_threshold_gated_case(self, case)
|
||||
|
||||
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:
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_dual_chunk_cuda_graph_decode_case(self, case)
|
||||
|
||||
# Layout-robustness. See dense/test_triton.py for the rationale.
|
||||
# dual_chunk_flash_attn EXTEND fails on non_monotonic_extend with
|
||||
# ~67% mismatch and max abs diff ~1.1. The dual-chunk prefill path
|
||||
# uses `cu_seqlens_*` indexing into a contiguous K layout
|
||||
# (see `_dual_chunk_flash_attn_prefill_func` in
|
||||
# dual_chunk_flashattention_backend.py:834+), which assumes K for
|
||||
# the new extend tokens is laid out contiguously in
|
||||
# `[begin, end)` slot order. Scattering extend-token slots within a
|
||||
# request breaks that contiguity. Documented as a known production
|
||||
# limitation.
|
||||
LAYOUT_ROBUSTNESS_CASES = (
|
||||
DualChunkAttentionCase(
|
||||
name="layout_dual_chunk_extend_two_request",
|
||||
backend="dual_chunk_flash_attn",
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
num_heads=4,
|
||||
num_kv_heads=4,
|
||||
page_size=16,
|
||||
prefix_lens=(0, 0),
|
||||
extend_lens=(16, 32),
|
||||
),
|
||||
DualChunkAttentionCase(
|
||||
name="layout_dual_chunk_decode_page_boundary",
|
||||
backend="dual_chunk_flash_attn",
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_heads=4,
|
||||
num_kv_heads=4,
|
||||
page_size=16,
|
||||
prefix_lens=(14, 15, 16),
|
||||
),
|
||||
)
|
||||
LAYOUT_KNOWN_FAILURES = {
|
||||
("layout_dual_chunk_extend_two_request", "non_monotonic_extend"): (
|
||||
"dual_chunk_flash_attn prefill uses cu_seqlens_* indexing "
|
||||
"into contiguous K slots within an extend "
|
||||
"(`_dual_chunk_flash_attn_prefill_func` in "
|
||||
"dual_chunk_flashattention_backend.py:834+); scattered "
|
||||
"extend-token slots break that contiguity."
|
||||
),
|
||||
}
|
||||
|
||||
def test_layout_robustness_cases(self):
|
||||
for case in self.LAYOUT_ROBUSTNESS_CASES:
|
||||
for layout in ("interleaved_pages", "non_monotonic_extend"):
|
||||
if layout == "non_monotonic_extend" and case.forward_mode.is_decode():
|
||||
continue
|
||||
reason = self.LAYOUT_KNOWN_FAILURES.get((case.name, layout))
|
||||
if reason is not None:
|
||||
print(
|
||||
f"[layout-known-failure] {case.name} x {layout}: {reason}",
|
||||
flush=True,
|
||||
)
|
||||
continue
|
||||
with self.subTest(case=case.name, layout=layout):
|
||||
run_dual_chunk_attention_case(self, case, loc_layout=layout)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -20,7 +20,6 @@ Columns are runner modes; rows are attention backends. Cells use:
|
||||
| `triton` | ✓ 10 input layouts (page 1/16/32, prefix/decode edges) | ✓ MLA decode page-boundary | ✓ ragged page-boundary extend | ✓ ragged page-boundary extend | ✓ EAGLE chain (topk=1) | ✓ EAGLE tree (topk=2) | — (V1 DE not enabled for Triton MLA; Triton uses V2 path) | — | ✓ fixed-tokens-per-req | ✓ chain (topk=1) + tree (topk=2) | ✓ via `DRAFT_EXTEND_V2` graph runner | — (no FKVMTP wiring for MLA) |
|
||||
| `flashinfer` | ✓ 10 input layouts with DeepSeek-like `kv_lora_rank=512`, `qk_rope_head_dim=64` | ✓ MLA decode page-boundary | ✓ ragged page-boundary extend | ✓ ragged page-boundary extend | ✓ EAGLE chain (topk=1) | ✓ EAGLE chain (topk=1) | ✓ EAGLE ragged-accept | ✓ EAGLE ragged-accept | blocked: `is_draft_extend()` default `include_v2=False` (`flashinfer_mla_backend.py:432,501,454-455,512`) | ✓ chain (topk=1) only — tree blocked by `topk=1` reject (`flashinfer_mla_backend.py:910-913`) | ✓ EAGLE ragged-accept (V1) | — (no FKVMTP wiring for MLA) |
|
||||
| `flashmla` | ✓ FlashMLA-compatible page-size-64 cases (zero-prefix exact page, input page edges 63/64/65, prefix exact page, total exact page, cross page, ragged, decode page-boundary, decode bsz=1 nonzero prefix) | ✓ page-size-64 decode page-boundary | ✓ ragged page-boundary extend | ✓ ragged page-boundary extend | ✓ EAGLE chain (topk=1) | ✓ EAGLE chain (topk=1) | ✓ EAGLE ragged-accept | deferred: parent FlashInfer-MLA capture path expects 1D `cuda_graph_kv_indices`, FlashMLA allocates 2D `[max_bs, (max_context+PAGE_SIZE)//PAGE_SIZE]` (`flashmla_backend.py:347-348` + parent `init_forward_metadata_capture_cuda_graph`) | — (FlashMLA does not implement V2) | ✓ chain (topk=1) only — tree blocked by `topk=1` reject (`flashmla_backend.py:555-558`) | — (DE CG deferred above) | — |
|
||||
| `cutlass_mla` | skip:hw — needs SM 10.0+ (Blackwell); current 1 case uses `ForwardMode.EXTEND` but `CutlassMLABackend` only overrides `forward_decode` (`cutlass_mla_backend.py:226`) and falls through to FlashInfer MLA for other modes → **case should be DECODE**; PAGE_SIZE fixed at 128 (`cutlass_mla_backend.py:31`) | — (decode-only backend; no extend/CG) | — | — | blocked: tree via `topk=1` reject inherited from FlashInfer MLA parent | — | — | — | — | — | — | — |
|
||||
| `trtllm_mla` | skip:hw — needs SM 12.0a / 12.1a (`is_sm120_supported`) | — | — | — | blocked: `topk=1` only (`trtllm_mla_backend.py:1223-1229` inherits from FlashInfer MLA) | — | — | — | — | — | — | — |
|
||||
| `tokenspeed_mla` | skip:hw — needs `find_spec("tokenspeed_mla")`, SM 10.0+, and `kv_cache_dtype=fp8_e4m3` (`server_args.py:2814-2818`); current MLA fixture does not emit FP8 KV cache | — | — | — | blocked: `topk=1` only (`tokenspeed_mla_backend.py:341-347` inherits from TRT-LLM MLA) | — | — | — | — | — | — | — |
|
||||
|
||||
@@ -60,17 +59,12 @@ multi-step draft backends and cannot ever appear at runtime.
|
||||
- **Tokenspeed MLA tree verify / draft-extend with `topk > 1`** —
|
||||
`TokenspeedMLAMultiStepDraftBackend` inherits from
|
||||
`TRTLLMMLAMultiStepDraftBackend` (`tokenspeed_mla_backend.py:341-347`).
|
||||
- **Cutlass MLA extend / verify / draft-extend** — `CutlassMLABackend` only
|
||||
overrides `forward_decode` (`cutlass_mla_backend.py:226`) and only handles
|
||||
`is_decode_or_idle` in `init_forward_metadata*` (`cutlass_mla_backend.py:86,
|
||||
156, 197`). Anything else falls through to FlashInfer MLA.
|
||||
- **FlashInfer-MLA `DRAFT_EXTEND_V2` graph capture/replay** —
|
||||
`flashinfer_mla_backend.py:432,501` only route through `is_draft_extend()`
|
||||
(default `include_v2=False`); `else: raise ValueError("Invalid mode")` at
|
||||
`flashinfer_mla_backend.py:454-455,512`.
|
||||
- **All MLA backends fixed page size** — FlashMLA forces `page_size=64`,
|
||||
Cutlass MLA forces `page_size=128`, TRT-LLM MLA and Tokenspeed MLA force
|
||||
`page_size in {32, 64}`.
|
||||
TRT-LLM MLA and Tokenspeed MLA force `page_size in {32, 64}`.
|
||||
|
||||
## Backend Container Gate (SM10.x)
|
||||
|
||||
@@ -91,9 +85,6 @@ See `KNOWN_FAILURES.md` §3 for the full root cause + fix.
|
||||
override capture/replay in `FlashMLABackend` to use its 2D layout, or
|
||||
allocate both parent-style 1D and FlashMLA-style 2D buffers and route
|
||||
`DRAFT_EXTEND` to the parent path).
|
||||
- Switch `mla/test_cutlass_mla.py` to `ForwardMode.DECODE` so it actually
|
||||
exercises `CutlassMLABackend.forward_decode` instead of falling through to
|
||||
FlashInfer MLA when SM 10.0+ is available.
|
||||
- Add hardware-gated tests for `cutlass_mla`, `trtllm_mla`, and `tokenspeed_mla`
|
||||
- Add hardware-gated tests for `trtllm_mla` and `tokenspeed_mla`
|
||||
decode (chain spec only) when the appropriate hardware/KV dtype fixtures are
|
||||
available.
|
||||
|
||||
@@ -1,98 +0,0 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
|
||||
MLAAttentionCase,
|
||||
run_mla_attention_case,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# Cutlass MLA requires exactly Blackwell SM 10.0. The sgl-kernel
|
||||
# `cutlass_mla_decode` checks `sm_version == 100` (major*10+minor), so
|
||||
# SM 10.3 (GB300) reports sm_version=103 and is rejected by the kernel.
|
||||
# PAGE_SIZE is fixed to 128 in the backend.
|
||||
_REQUIRED_SM_MAJOR = 10
|
||||
_REQUIRED_SM_MINOR = 0
|
||||
|
||||
MLA_SHAPE_KWARGS = dict(
|
||||
kv_lora_rank=512,
|
||||
qk_rope_head_dim=64,
|
||||
hidden_size=1024,
|
||||
max_context_len=256,
|
||||
)
|
||||
|
||||
|
||||
def _supported() -> tuple[bool, str]:
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA is required"
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
if major != _REQUIRED_SM_MAJOR or minor != _REQUIRED_SM_MINOR:
|
||||
return (
|
||||
False,
|
||||
f"cutlass_mla requires exactly SM {_REQUIRED_SM_MAJOR}.{_REQUIRED_SM_MINOR} "
|
||||
f"(B200 Blackwell); got SM {major}.{minor}",
|
||||
)
|
||||
return True, ""
|
||||
|
||||
|
||||
_SUPPORTED, _SKIP_REASON = _supported()
|
||||
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=12, stage="base-b", runner_config="4-gpu-b200")
|
||||
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
@unittest.skipIf(not _SUPPORTED, _SKIP_REASON)
|
||||
class TestCutlassMLAAttentionBackendCorrectness(CustomTestCase):
|
||||
# CutlassMLABackend only overrides `forward_decode`; EXTEND falls through
|
||||
# to the FlashInferMLAAttnBackend parent and bypasses cutlass code
|
||||
# entirely. Use DECODE so the test actually exercises the cutlass kernel
|
||||
# on Blackwell. Page size is fixed to PAGE_SIZE=128 (server_args.py
|
||||
# forces this for cutlass_mla).
|
||||
CASES = (
|
||||
MLAAttentionCase(
|
||||
name="mla_decode_cutlass_page_boundary",
|
||||
backend="cutlass_mla",
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_heads=4,
|
||||
page_size=128,
|
||||
prefix_lens=(126, 127, 128),
|
||||
),
|
||||
MLAAttentionCase(
|
||||
name="mla_decode_cutlass_bsz1_nonzero_prefix",
|
||||
backend="cutlass_mla",
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_heads=4,
|
||||
page_size=128,
|
||||
prefix_lens=(63,),
|
||||
),
|
||||
MLAAttentionCase(
|
||||
name="mla_decode_cutlass_above_page",
|
||||
backend="cutlass_mla",
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_heads=4,
|
||||
page_size=128,
|
||||
prefix_lens=(128, 129, 130),
|
||||
),
|
||||
MLAAttentionCase(
|
||||
name="mla_decode_cutlass_multi_page",
|
||||
backend="cutlass_mla",
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_heads=4,
|
||||
page_size=128,
|
||||
prefix_lens=(127, 200, 255),
|
||||
),
|
||||
)
|
||||
|
||||
def test_projected_mla_attention_cases(self):
|
||||
for case in self.CASES:
|
||||
with self.subTest(case=case.name, backend=case.backend):
|
||||
run_mla_attention_case(self, case, **MLA_SHAPE_KWARGS)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -15,13 +15,6 @@ register_cuda_ci(
|
||||
est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize as aot_awq_dequantize
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
if IS_CI:
|
||||
@@ -34,59 +27,16 @@ else:
|
||||
configs = list(itertools.product(qweight_row_range, qweight_cols_range))
|
||||
|
||||
|
||||
def check_correctness():
|
||||
if not AOT_AVAILABLE:
|
||||
print("sgl_kernel AOT not available, skipping correctness check")
|
||||
return
|
||||
|
||||
qweight_row, qweight_col = 128, 16
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
|
||||
print("Correctness check passed (JIT vs AOT)")
|
||||
|
||||
|
||||
if AOT_AVAILABLE:
|
||||
line_vals = ["jit", "aot"]
|
||||
line_names = ["JIT Kernel", "AOT Kernel"]
|
||||
styles = [("blue", "-"), ("green", "-")]
|
||||
else:
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT Kernel"]
|
||||
styles = [("blue", "-")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["qweight_row", "qweight_col"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
line_vals=["jit"],
|
||||
line_names=["JIT Kernel"],
|
||||
styles=[("blue", "-")],
|
||||
ylabel="us",
|
||||
plot_name="awq-dequantize-jit-vs-aot",
|
||||
plot_name="awq-dequantize-jit",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
@@ -111,16 +61,12 @@ def benchmark(qweight_row, qweight_col, provider):
|
||||
device=device,
|
||||
)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
|
||||
elif provider == "aot":
|
||||
fn = lambda: aot_awq_dequantize(qweight, scales, qzeros)
|
||||
else:
|
||||
if provider != "jit":
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
check_correctness()
|
||||
benchmark.run(print_data=True)
|
||||
|
||||
@@ -11,13 +11,6 @@ from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize as aot_awq_dequantize
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
|
||||
def reverse_awq_order(t: torch.Tensor):
|
||||
bits = 4
|
||||
@@ -117,54 +110,5 @@ def test_awq_dequantize_jit_vs_torch(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"qweight_row,qweight_col,is_bf16_act",
|
||||
list(
|
||||
itertools.product(
|
||||
[128, 256, 512, 1024, 3584],
|
||||
[16, 32, 64, 128, 448],
|
||||
[True, False],
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_awq_dequantize_jit_vs_aot(
|
||||
qweight_row: int, qweight_col: int, is_bf16_act: bool
|
||||
):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
|
||||
if is_bf16_act:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
|
||||
else:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Run both implementations
|
||||
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
|
||||
@@ -1,60 +0,0 @@
|
||||
"""GPTQ builds its per-layer scheme lazily in `create_weights`, so the layer has
|
||||
to declare `scheme = None` for the `is None` probe to see it.
|
||||
|
||||
Regression: the probe used to be `hasattr(layer, "scheme")`, which degraded to
|
||||
always-true once `LinearBase` grew that class default -- the scheme was never
|
||||
built and every GPTQ model died with ``'NoneType' object has no attribute
|
||||
'create_weights'``.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.linear import LinearBase, ReplicatedLinear
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig
|
||||
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
||||
|
||||
_GPTQ_CHECKPOINT_CONFIG = {
|
||||
"bits": 4,
|
||||
"group_size": 128,
|
||||
"desc_act": False,
|
||||
"lm_head": False,
|
||||
"dynamic": {},
|
||||
"checkpoint_format": "gptq",
|
||||
"true_sequential": True,
|
||||
"static_groups": False,
|
||||
}
|
||||
|
||||
|
||||
class TestGPTQSchemeAttach(CustomTestCase):
|
||||
def test_linear_layer_gets_a_scheme(self):
|
||||
layer = ReplicatedLinear(
|
||||
input_size=256,
|
||||
output_size=128,
|
||||
bias=False,
|
||||
params_dtype=torch.float16,
|
||||
quant_config=GPTQConfig.from_config(_GPTQ_CHECKPOINT_CONFIG),
|
||||
prefix="model.layers.0.mlp.down_proj",
|
||||
)
|
||||
self.assertIsNotNone(layer.scheme)
|
||||
self.assertTrue(hasattr(layer, "qweight"))
|
||||
|
||||
def test_scheme_default_is_declared_on_every_quantizable_layer_base(self):
|
||||
"""`get_linear_quant_method` hands a linear method a `LinearBase` or a
|
||||
quantized `ParallelLMHead`; `GPTQMarlinConfig` hands
|
||||
`GPTQMarlinMoEMethod` a bare `FusedMoE`. The MoE attach has no e2e
|
||||
coverage, so this is its only guard.
|
||||
"""
|
||||
self.assertIsNone(LinearBase.scheme)
|
||||
self.assertIsNone(VocabParallelEmbedding.scheme)
|
||||
self.assertIsNone(FusedMoE.scheme)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -101,7 +101,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
|
||||
# MLA-family kernels that must never leak into the MHA arm.
|
||||
MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla")
|
||||
# No kernel-facing-id wiring anywhere: must stay rejected until they get one.
|
||||
UNWIRED_BACKENDS = ("cutlass_mla", "aiter")
|
||||
UNWIRED_BACKENDS = ("aiter",)
|
||||
|
||||
def test_triton_allowed_on_every_arm(self):
|
||||
"""Triton reads both view families, so it is the one backend neither
|
||||
|
||||
@@ -1423,7 +1423,6 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
from sglang.srt.arg_groups.overrides import (
|
||||
ResolvedView,
|
||||
_attention_backend_default,
|
||||
_attention_backend_dual_chunk,
|
||||
_attention_backend_fa3_fp8_fallback,
|
||||
_attention_backend_platform_fallbacks,
|
||||
)
|
||||
@@ -1459,19 +1458,6 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
with override_platform(has_amx=True):
|
||||
self.assertEqual(_attention_backend_platform_fallbacks(view), {})
|
||||
|
||||
# dual-chunk config: mismatched explicit backend raises verbatim
|
||||
def _mc(dual):
|
||||
return SimpleNamespace(
|
||||
_model_config=SimpleNamespace(
|
||||
hf_config=SimpleNamespace(dual_chunk_attention_config=dual)
|
||||
),
|
||||
attention_backend="fa3",
|
||||
)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
_attention_backend_dual_chunk(ResolvedView(_mc({"a": 1})))
|
||||
self.assertEqual(_attention_backend_dual_chunk(ResolvedView(_mc(None))), {})
|
||||
|
||||
def test_dllm_platform_paths_at_callable_level(self):
|
||||
from sglang.srt.arg_groups.overrides import (
|
||||
ResolvedView,
|
||||
@@ -2668,16 +2654,6 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
),
|
||||
{"page_size": 64},
|
||||
)
|
||||
# chained: cutlass_mla decode -> 128, then trtllm_mha prefill keeps 128
|
||||
self.assertEqual(
|
||||
_mla_backend_page_constraints(
|
||||
_view(
|
||||
decode_attention_backend="cutlass_mla",
|
||||
prefill_attention_backend="trtllm_mha",
|
||||
)
|
||||
),
|
||||
{"page_size": 128},
|
||||
)
|
||||
# no matching backend: nothing declared
|
||||
self.assertEqual(_mla_backend_page_constraints(_view()), {})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user