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
2026-09-10 15:12:01 +08:00
committed by GitHub
co-authored by Brayden Zhong
parent 7152c14384
commit c0b790cf7f
71 changed files with 48 additions and 12654 deletions
@@ -1190,8 +1190,8 @@ def dsa_impl_capability(impl: str) -> tuple[bool, str]:
# TRT-LLM Gen FMHA / MLA require Blackwell SM10.0 (B200 NVL).
# SM10.3 (GB300) raises "Missing TRTLLM-GEN kernel" at runtime because
# the kernel binary in the container isn't compiled for sm_103.
# Require exactly SM10.0 (same constraint as cutlass_mla) until the
# container ships sm_103-compiled TRTLLM-GEN kernels.
# Require exactly SM10.0 until the container ships sm_103-compiled
# TRTLLM-GEN kernels.
if major != 10 or minor != 0:
return (
False,
@@ -73,20 +73,6 @@ from ..attention_methods.dsv4_attention import (
run_dsv4_fixture_eager,
run_dsv4_forward,
)
from ..attention_methods.dual_chunk_attention import (
DualChunkAttentionCase,
_clone_dual_chunk_cache,
_restore_dual_chunk_cache,
build_dual_chunk_attention_fixture,
dual_chunk_fixture_inputs,
expected_dual_chunk_output_from_inputs,
make_dual_chunk_case_with_prefix_lens,
make_dual_chunk_random_inputs,
make_dual_chunk_replay_inputs,
prepare_dual_chunk_runner_inputs,
run_dual_chunk_fixture_eager,
run_dual_chunk_forward,
)
from ..attention_methods.gdn_attention import DEFAULT_DEVICE as GDN_DEFAULT_DEVICE
from ..attention_methods.gdn_attention import DEFAULT_DTYPE as GDN_DEFAULT_DTYPE
from ..attention_methods.gdn_attention import (
@@ -891,58 +877,3 @@ def run_dsa_sparse_cuda_graph_decode_case(
dtype=dtype,
device=device,
)
def run_dual_chunk_cuda_graph_decode_case(
testcase,
case: DualChunkAttentionCase,
*,
head_dim: int = DEFAULT_HEAD_DIM,
hidden_size: int = DEFAULT_HIDDEN_SIZE,
max_context_len: int = DENSE_DEFAULT_MAX_CONTEXT_LEN,
dtype: torch.dtype = DENSE_DEFAULT_DTYPE,
device: str = DENSE_DEFAULT_DEVICE,
cuda_graph_capture_batch_size: int | None = None,
):
"""Dual-chunk CUDA-graph decode replay. Decode reads cached K/V (set
by `set_kv_buffer` inside `forward_decode`) so the capture/replay
contract is the same shape as dense attention. The
`_clone_dual_chunk_cache` / `_restore_dual_chunk_cache` hooks snapshot
both K and V buffers so the capture forward's writes don't bleed into
replay state."""
if not case.forward_mode.is_decode():
raise ValueError("run_dual_chunk_cuda_graph_decode_case expects a DECODE case.")
capture_batch_size = cuda_graph_capture_batch_size or case.batch_size
adapter = CudaGraphDecodeAdapter(
build_fixture=build_dual_chunk_attention_fixture,
make_case=make_dual_chunk_case_with_prefix_lens,
make_forward_batch=_make_dense_forward_batch,
fixture_inputs=dual_chunk_fixture_inputs,
make_capture_inputs=make_dual_chunk_random_inputs,
make_replay_inputs=make_dual_chunk_replay_inputs,
prepare_inputs=prepare_dual_chunk_runner_inputs,
run_eager=run_dual_chunk_fixture_eager,
run_forward=run_dual_chunk_forward,
expected_output=expected_dual_chunk_output_from_inputs,
clone_state=_clone_dual_chunk_cache,
restore_state=_restore_dual_chunk_cache,
allow_padding=True,
atol=DENSE_ATOL,
rtol=DENSE_RTOL,
)
_run_cuda_graph_decode_case(
testcase,
case,
adapter=adapter,
build_kwargs=dict(
head_dim=head_dim,
hidden_size=hidden_size,
max_context_len=max_context_len,
dtype=dtype,
device=device,
),
capture_batch_size=capture_batch_size,
max_context_len=max_context_len,
dtype=dtype,
device=device,
)