From a076ec1a7af0aa12ee1260a8821e4e7992c2c840 Mon Sep 17 00:00:00 2001
From: b8zhong <b8zhong@uwaterloo.ca>
Date: Thu, 30 Oct 2025 22:33:37 -0700
Subject: [PATCH] Revert "fix llama4 kv cache layout" (#12437)

---
 docs/advanced_features/attention_backend.md | 2 +-
 python/sglang/srt/server_args.py            | 7 -------
 2 files changed, 1 insertion(+), 8 deletions(-)

diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md
index ec223add1..bbe39e0a1 100644
--- a/docs/advanced_features/attention_backend.md
+++ b/docs/advanced_features/attention_backend.md
@@ -21,7 +21,7 @@ The support matrix is split into two parts: MHA (standard attention) and MLA (mu
 | **Triton**                      | ❌                          | ❌               | ✅              | ✅              | ✅                 | ✅             |
 | **Torch Native (SDPA)**         | ❌                          | ❌               | ❌              | ❌              | ❌                 | ❌             |
 | **FlexAttention (PyTorch)**     | ❌                          | ❌               | ❌              | ❌              | ❌                 | ❌             |
-| **TRTLLM MHA**                  | 16, 32 or 64                | ❌               | ✅              | ❌              | ❌                 | ❌             |
+| **TRTLLM MHA**                  | 16, 32 or 64                | ✅               | ✅              | ❌              | ❌                 | ❌             |
 | **Dual Chunk FlashAttention**   | ✅                          | ❌               | ❌              | ❌              | ❌                 | ❌             |
 | **AITER (ROCm)**                | ✅                          | ❌               | ✅              | ✅              | ❌                 | ❌             |
 | **Wave (ROCm)**                 | ✅                          | ❌               | ❌              | ❌              | ❌                 | ❌             |
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 2eadb19d6..3f130e3ca 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -980,13 +980,6 @@ class ServerArgs:
                 logger.warning(
                     "Use trtllm_mha as attention backend on sm100 for Llama4 model"
                 )
-            if is_sm100_supported() and self.attention_backend == "trtllm_mha":
-                # TODO(brayden): remove this once TRTLLM MHA kernel for FP8 w/ tileSizeKv=128 is available.
-                # This is a Llama 4 specific issue only.
-                self.kv_cache_dtype = "bfloat16"
-                logger.warning(
-                    "Setting kv_cache_dtype to bfloat16 for Llama4 with trtllm_mha backend, due to a missing FlashInfer TRTLLM MHA kernel for FP8 KV Cache"
-                )
             if is_sm100_supported() and self.moe_runner_backend == "auto":
                 if self.quantization in {"fp8", "modelopt_fp8"}:
                     self.moe_runner_backend = "flashinfer_trtllm"