Why Attention Architecture Now Determines Inference Speed

Agentic and long-context workloads are no longer edge cases—they are the default. As context lengths grow, attention consumes an ever-larger share of inference time. The way you design attention (not just implement it) now directly impacts model performance on GPUs. This is the core premise of AI model co-design: shaping model architecture around how GPUs actually execute it.

In this deep dive, we analyze how three key parameters—group size (G), head dimension (Hsz), and sequence length (ISL/KVSL)—affect dense attention performance. We ground every insight in analytical formulas (GEMM-shape arithmetic) and measured data from prefill and decode kernels using FP8 for both attention compute and KV cache. For a broader discussion on hardware-friendly LLM design, see our previous analysis on AI model co-design.

Before diving in, let's clarify the two phases:

  • Prefill: Processes the full prompt in parallel → large GEMM-M (ISL × G), compute-bound.
  • Decode: Generates one token at a time → small GEMM-M (G), memory-bound by KV cache reads from HBM.

Speculative decoding can shift decode toward compute-bound by increasing effective GEMM-M. Each parameter below is analyzed separately for prefill and decode because their bottlenecks differ.

Diagram illustrating attention head grouping (GQA) and its impact on GPU memory bandwidth during LLM inference Developer Related Image

The Four Co-Design Guidelines (with Code & Math)

Guideline 1: Choose Group Size (G) for Decode Efficiency—and Push It High

Group size (G) is the number of query heads sharing one KV head. MHA has G=1, GQA has G=4,8,16…, MQA has G=QH.

Prefill is dominated by ISL, not G. Arithmetic intensity ≈ 2 × ISL / (1 + 1/G). At ISL=32K, raising G from 8 to 16 improves AI by <6%. Runtime changes by <1% across G=1 to 64.

Decode arithmetic intensity ≈ 2 × G (when KVSL ≫ G). Doubling G doubles decode AI and roughly halves runtime (until overheads dominate).

# Quick calculation: decode arithmetic intensity vs group size
# Assumes FP8 (1 byte per element)
def decode_arithmetic_intensity(G, KVSL):
    # FLOPs = 4 * DB * QH * KVSL * Hsz (constant in G)
    # Bytes = 2 * DB * KH * Hsz * (G + KVSL)
    # AI = 2 * G * KVSL / (G + KVSL)
    return 2 * G * KVSL / (G + KVSL)

for G in [1, 8, 16, 32]:
    ai = decode_arithmetic_intensity(G, 32768)
    print(f"G={G:2d} -> AI ≈ {ai:.1f} FLOPs/byte")
# Output: G=1 -> 2.0, G=8 -> 16.0, G=16 -> 32.0, G=32 -> 64.0

Takeaway: Higher G improves decode speed and GPU utilization. Speculative decoding is another lever.

Guideline 2: Use Head Dimension (Hsz) = 128 or 256

Hsz does not change arithmetic intensity (doubles both FLOPs and bytes). But runtime increases with Hsz because:

  • Matmul grows in aligned steps; Hsz=64 still pays for a 128-wide tile. Hsz≥512 hits TMEM limits.
  • Memory moves in 128-byte units; Hsz should be multiple of 128.
  • Softmax cost is independent of Hsz (operates on score matrix).
# Softmax ops scale with sequence length, not head dimension
def softmax_ops_prefill(PB, QH, ISL):
    return PB * QH * ISL * ISL  # independent of Hsz

def softmax_ops_decode(DB, QH, KVSL):
    return DB * QH * KVSL      # independent of Hsz

Takeaway: Hsz=128 or 256 is the sweet spot for GPU alignment and TMEM budget.

Guideline 3: Reduce Effective KV State Where Possible

Prefill scales quadratically with ISL (every token attends to every token), while decode scales linearly with KVSL (each step reads full KV cache). The cost is asymmetric.

Strategies: KV-cache compression, sparse/sliding-window attention, hybrid architectures (e.g., only some layers carry global KV state).

Guideline 4: Let KH Determine Parallelism Strategy

Tensor parallelism (TP) splits attention heads across GPUs. Practical limit: TP ≤ KH so each GPU owns at least one complete group (one KV head + its G query heads). Once TP > KH, KV state is duplicated across ranks, wasting memory and bandwidth.

For models with few KV heads (MQA with 1, GQA with 2), TP exhausts quickly. Use Attention Data Parallelism (ADP) or KV Parallelism (KVP) for attention, plus Expert Parallelism (EP) for MoE FFN. TensorRT-LLM implements these as Wide EP (ADP+EP) and Helix Parallelism (KVP+EP).

# Parallelism check: ensure TP does not exceed KV head count
def validate_tp(KH, TP):
    if TP > KH:
        raise ValueError(f"TP={TP} exceeds KH={KH}. KV state will be duplicated.")
    return True

# Example: GQA with 2 KV heads
validate_tp(KH=2, TP=2)  # OK
# validate_tp(KH=2, TP=4)  # Raises error

GPU server rack with high-bandwidth memory handling long-context KV cache for AI model decoding Dev Environment Setup

Critical Caveats and Limitations

  • Prefill vs decode are not symmetric. Optimizations that help decode (higher G) barely move prefill. Always profile both phases separately.
  • Speculative decoding changes the game. It raises effective GEMM-M to (1+D)×G, where D is draft tokens. Once large enough to fill compute tiles, decode shifts toward compute-bound—reducing the relative benefit of high G.
  • Overheads dominate at short sequences. At KVSL=32K, decode runtime flattens beyond G=16 due to fixed setup, post-processing, and flash-decoding reduction overhead. Longer contexts (128K) better amortize these costs.
  • TP ≤ KH is a hard limit. Violating it duplicates KV state and negates scaling benefits. Models with few KV heads must adopt ADP/KVP + EP instead.
  • Hsz=64 is inefficient. It still pays for a 128-wide tile—you get half the compute for the same cost. Avoid unless constrained.

Next Steps for Learning

  • Experiment with GQA group sizes in your own model: measure decode throughput at G=1, 4, 8, 16 on your target GPU.
  • Profile prefill and decode separately using NVIDIA Nsight Systems to identify whether you are compute-bound or memory-bound.
  • Explore TensorRT-LLM's Wide EP and Helix Parallelism for scaling models with few KV heads.
  • Read the companion post on sparse attention (coming soon) to learn how sparsity changes the arithmetic intensity calculus.
  • For a related perspective on enterprise AI integration, check out our analysis on how AI agents are redefining enterprise integration.

Performance graph showing prefill vs decode runtime scaling with sequence length on NVIDIA GPUs System Abstract Visual

Conclusion: Co-Design Is Not Optional

The four guidelines above form a practical checklist for model developers targeting NVIDIA GPUs:

  1. Push G high for decode efficiency (prefill is insensitive).
  2. Use Hsz=128 or 256 to align with GPU tiles and 128-byte transfers.
  3. Reduce effective KV state via compression, sparsity, or hybrid architectures.
  4. Match parallelism to KH—keep TP ≤ KH and use Wide EP/Helix for few-KV-head models.

As agentic workloads become standard, attention architecture is no longer an afterthought—it is the primary lever for inference performance. Co-designing model architecture with GPU execution in mind is how you deliver fast, interactive long-context inference at scale.

Stay secure: As you integrate AI agents into your stack, be aware of new supply chain risks like AGENTS.md injection attacks. Security must be co-designed too.

This content was drafted using AI tools based on reliable sources, and has been reviewed by our editorial team before publication. It is not intended to replace professional advice.